# Graph-based Neural Acceleration
for Nonnegative Matrix Factorization

Jens Sjölund  Maria Bånkestad

###### Abstract

We describe a graph-based neural acceleration technique for nonnegative matrix
factorization that builds upon a connection between matrices and bipartite
graphs that is well-known in certain fields, e.g., sparse linear algebra, but
has not yet been exploited to design graph neural networks for matrix
computations. We first consider low-rank factorization more broadly and
propose a graph representation of the problem suited for graph neural
networks. Then, we focus on the task of nonnegative matrix factorization and
propose a graph neural network that interleaves bipartite self-attention
layers with updates based on the alternating direction method of multipliers.
Our empirical evaluation on synthetic and two real-world datasets shows that
we attain substantial acceleration, even though we only train in an
unsupervised fashion on smaller synthetic instances.

Graph Neural Network, Data-driven optimization, Nonnegative Matrix
Factorization

##  1 Introduction

Nonnegative matrix factorization (NMF) is a standard tool for analyzing high-
dimensional nonnegative data such as images, text, and audio spectra (Fu et
al., 2019; Gillis, 2020). In NMF, the aim is to find an approximate low-rank
factorization of a given data matrix ${\bm{V}}\in\mathbb{R}^{m\times n}$
according to

${\bm{V}}\approx{\bm{W}}{\bm{H}}^{\top},\quad{\bm{W}}\geq 0,{\bm{H}}\geq 0,$
(1)

where the basis matrix ${\bm{W}}\in\mathbb{R}^{m\times r}$ and the mixture
matrix ${\bm{H}}\in\mathbb{R}^{n\times r}$ are elementwise nonnegative
matrices of rank $r\leq\min(m,n)$.

The nonnegativity constraints are the source of both the interpretability and
the complexity of NMF. In their absence, the optimal solution can be computed
efficiently using e.g., a singular value decomposition (SVD). In their
presence, the problem becomes NP-hard (Vavasis, 2010). State-of-the-art
algorithms for general NMF are essentially direct applications of standard
optimization techniques (Huang et al., 2016; Gillis, 2020), which are
sensitive to the initialization and may require many iterations to converge.

Like other data-driven optimization methods (Li & Malik, 2016; Chen et al.,
2021), we exploit that this problem has a specific structure, which we can
_learn_ to take advantage of. However, a practically useful NMF method must be
able to factorize matrices in a wide range of sizes, which means that—unlike
most other applications where data-driven optimization has been used—our
method must be able to solve problems of _variable_ size.

To the best of our knowledge, such methods have previously only been explored
for combinatorial optimization problems, which are often stated directly in
terms of graphs and are therefore well-suited to graph neural networks (Gasse
et al., 2019; Cappart et al., 2021). There is, however, a direct connection
also between matrices and graphs. This connection is well-known in certain
fields, e.g. sparse linear algebra (Kepner & Gilbert, 2011; Duff et al.,
2017), but has not been exploited to design graph neural networks that learn
to perform matrix computations—until now.

We describe a graph-based neural acceleration technique for nonnegative matrix
factorization. Our main contributions are as follows, we (i) formulate low-
rank matrix factorization in terms of a weighted bipartite graph; (ii) propose
a graph neural network for alternating minimization, which interleaves
Transformer-like updates with updates based on the alternating direction
method of multipliers (ADMM); (iii) demonstrate substantial acceleration on
synthetic matrices, also when applied to larger matrices than trained on; and
(iv) modify the synthetic data generation to achieve successful generalization
to two real-world datasets.

##  2 Related Work

Nonnegative matrix factorization has a wide range of applications, for
instance in image processing (Fu et al., 2019), bioinformatics (Kim & Park,
2007), and document analysis (Shahnaz et al., 2006), in part because it tends
to produce interpretable results (Lee & Seung, 1999). Since the problem is NP-
hard (Vavasis, 2010), this has led to the development of several algorithms
specialized for NMF. A few of them exploit the special analytical structure of
the problem (Recht et al., 2012; Gillis & Vavasis, 2013), but the majority are
adaptations of standard optimization methods (Lin, 2007; Kim & Park, 2008; Xu
& Yin, 2013; Gillis, 2020).

In contrast, learning to optimize is an emerging approach where the aim is to
learn optimization methods specialized to perform well on a set of training
problems (Chen et al., 2021). The same approach is also referred to as data-
driven optimization in more mathematically oriented works (Banert et al.,
2020), or as meta-learning or learning-to-learn when referring to the more
narrow scope of learning parameters for machine learning models (Li & Malik,
2016; Andrychowicz et al., 2016; Bello et al., 2017). It is related to
automated machine learning, “AutoML” (Hutter et al., 2019), but rather than
only selecting from or tuning a few existing optimizers, learning to optimize
generally involves a search over a much larger space of parameterized methods
that encompass entirely new optimization methods (Maheswaranathan et al.,
2020).

The iterative nature of optimization methods is often captured using either
feedforward or recurrent neural networks. In the feedforward approach, which
we use, the algorithm is unrolled over a limited number of iterations and
intertwined with learned components (Gregor & LeCun, 2010; Domke, 2012; Monga
et al., 2021). This makes it easier to train than the recurrent approach
(Andrychowicz et al., 2016; Venkataraman & Amos, 2021), but computational
considerations limit the number of iterations (depth) that can be considered.

Data-driven optimization methods applied to image-based problems (Diamond et
al., 2017; Banert et al., 2020) typically rely on convolutional neural
networks (CNNs). However, the geometry of matrices, when used in computations,
is decidedly different from images. Matrices have a non-Euclidean structure
and violate the assumption of translational equivariance that underlies
regular CNNs. This motivates the need to use tools from geometric deep
learning (Bronstein et al., 2017), and graph convolutional networks in
particular (Kipf & Welling, 2017; Gilmer et al., 2017).

For similar reasons, there has been a recent surge of interest in graph neural
networks for combinatorial optimization problems (Dai et al., 2017; Gasse et
al., 2019; Nair et al., 2020; Cappart et al., 2021), which are then often
represented as bipartite graphs between variables and constraints. Another
problem that is naturally represented as a bipartite graph is (low-rank)
matrix completion, which has, indeed, also been successfully tackled with
graph neural networks (Berg et al., 2017; Monti et al., 2017).

There is a vast literature building upon the connection between graphs and
matrices, both as an efficient way to implement graph algorithms (Kepner &
Gilbert, 2011) and as a theoretical tool for (sparse) linear algebra (Brualdi
& Cvetkovic, 2008; Duff et al., 2017). However, to the best of our knowledge,
no one has previously used it to design graph neural networks for numerical
linear algebra.

##  3 Background

In the literature on nonnegative matrix factorization, the most common way to
measure the quality of the approximation, $V\approx WH^{\top}$, is the
Frobenius norm, leading to the optimization problem

$\displaystyle\underset{{\bm{W}},{\bm{H}}}{\text{minimize}}$
$\displaystyle\frac{1}{2}\|{\bm{W}}{\bm{H}}^{\top}-{\bm{V}}\|_{\text{F}}^{2}$
(2) subject to $\displaystyle{\bm{W}},{\bm{H}}\geq 0.$

If either of the factors ${\bm{W}}$ or ${\bm{H}}$ is held fixed, what remains
is a nonnegative least-squares problem, which is convex. For this reason, most
NMF methods use an alternating optimization scheme (Lee & Seung, 2001; Xu &
Yin, 2013; Gillis, 2020). As these methods are—at best—guaranteed to converge
to a stationary point, it is important to initialize them appropriately.
Popular initialization methods include nonnegative double SVD (Boutsidis &
Gallopoulos, 2008) and (structured) random methods (Albright et al., 2006).
Our method provides a factorization that can be used by itself or as a refined
initialization that other methods can process further. Its underpinning is a
representation of matrix factorization as a bipartite graph.

###  3.1 Matrices as Graphs

Consider a graph, ${\mathcal{G}}=({\mathcal{V}},{\mathcal{E}})$, with nodes
$v_{i}\in{\mathcal{V}}$, and directed edges $e_{i,j}\in{\mathcal{E}}$ going
from a source node $v_{i}$ to a target node $v_{j}$. An undirected graph can
be modeled as having directed edges in both ways. A path is a sequence of
distinct nodes connected by directed edges.

Figure 1: A matrix corresponds to a weighted bipartite graph.

We may enrich the graph by associating features,
${\bm{x}}_{i}\in\mathbb{R}^{p}$ and ${\bm{e}}_{i,j}\in\mathbb{R}^{q}$, to the
nodes and edges, respectively.

The starting point for our method is the fact that any matrix
${\bm{V}}\in\mathbb{R}^{m\times n}$ can be represented as a directed,
weighted, and bipartite graph ${\mathcal{G}}({\bm{V}})$ called the König
digraph (Doob, 1984). The König digraph has $(m+n)$ nodes—$m$ row nodes and
$n$ column nodes—with edges going from row nodes to column nodes. The weight
of each edge is given by the corresponding element in the matrix,
${\bm{e}}_{i,j}={V}_{i,j}$, i.e. the edge features are scalar-valued ($q=1$).
An example is shown in Figure 1. By convention, edges corresponding to zeros
in the matrix are omitted. Transposing a matrix corresponds to reversing the
edges in the associated König digraph. The weight of a path is the product of
the weights along the path.

###  3.2 Message Passing Neural Networks

The recent surge of interest in graph neural networks has mainly focused on
message passing neural networks (Gilmer et al., 2017) which encompasses graph
convolutional networks (Kipf & Welling, 2017). In message passing neural
networks, a forward pass consists of multiple iterations of message passing,
optionally followed by a readout. Each message passing iteration $t$ consists
of evaluating message functions ${\bm{m}}^{t}_{i\to j}$ for every edge
${\bm{e}}_{i,j}$, aggregating these at each node using a permutation-invariant
aggregation function $\dottedsquare$ (typically the arithmetic mean or sum)
and updating the node features using an update function ${\bm{u}}^{t}$,

$\displaystyle{\bm{m}}^{t+1}_{j}$
$\displaystyle=\underset{i\in\mathcal{N}(j)}{\dottedsquare}{\bm{m}}^{t}_{i\to
j}\left({\bm{x}}_{i}^{t},{\bm{x}}_{j}^{t},{\bm{e}}_{i,j}\right),$ (3)
$\displaystyle{\bm{x}}^{t+1}_{j}$
$\displaystyle={\bm{u}}^{t}\left({\bm{x}}^{t}_{j},{\bm{m}}^{t+1}_{j}\right),$
(4)

where $\mathcal{N}(j)$ denotes the neighbours of node $v_{j}$. In the readout
phase, a readout function is used to compute a feature vector for the whole
graph, but since we work directly with the node features in this work we do
not use a readout.

##  4 Graph Representation of Constrained Low-rank Factorization

We will now describe our graph representation of the NMF problem, which is
generally applicable to constrained low-rank factorization. By representing
the problem in a format suitable for graph convolutional neural networks, we
make it possible to learn a fitting method tailored to the problem. In
particular, our representation of the problem as a bipartite graph makes it
possible to implement the graph neural network as a neural acceleration scheme
for conventional alternating optimization methods.

(13, 7)(0, 2)(4, 9)(11, 5)$W$$V$(6, 12)(8, 10)(3, 1)$H$ Figure 2: A simple
example of how we represent a matrix factorization
${\bm{V}}={\bm{W}}{\bm{H}}^{\top}$ as a bipartite graph. To reduce clutter,
the edge weights $V_{ij}$ are not shown.

Before going into the derivation, consider the following factorization

$\underbrace{\begin{pmatrix}162&174&46\\\ 24&20&2\\\ 132&122&21\\\
126&138&38\end{pmatrix}}_{\bm{V}}=\underbrace{\begin{pmatrix}13&7\\\ 0&2\\\
4&9\\\ 11&5\end{pmatrix}}_{\bm{W}}\underbrace{\begin{pmatrix}6&8&3\\\
12&10&1\end{pmatrix}}_{{\bm{H}}^{\top}},$ (5)

that we represent as a bipartite graph according to Figure 2.

Figure 3: We represent matrix multiplication as a graph by expressing it as a
sum of outer products, which are processed separately until the final,
concatenation, step. The processing consists of (i) representing each vector
as a König digraph, (ii) forming the concatenation graph, (iii) mapping each
concatenation graph to the corresponding line digraph, and (iv) concatenating
the node features of the line digraphs.

Matrix multiplication can be expressed in terms of König digraphs. Given
${\bm{W}}\in\mathbb{R}^{m\times r}$ and ${\bm{H}}\in\mathbb{R}^{n\times r}$,
we first define the concatenation graph
${\mathcal{G}}({\bm{W}})*{\mathcal{G}}({\bm{H}}^{\top})$ by identifying the
$r$ column nodes of ${\mathcal{G}}({\bm{W}})$ with the $r$ row nodes of
${\mathcal{G}}({\bm{H}}^{\top})$. Visually, this corresponds to “gluing
together” the digraphs. The $(i,j)$ entry of the matrix product
${\bm{W}}{\bm{H}}^{\top}$ is the sum of the weights of all paths in the
concatenation graph ${\mathcal{G}}({\bm{W}})*{\mathcal{G}}({\bm{H}}^{\top})$
from row node $i$ to column node $j$.

Our derivation, illustrated in Figure 3, starts with expressing matrix
multiplication as a sum of outer products,

${\bm{W}}{\bm{H}}^{\top}=\sum_{c=1}^{r}{\bm{w}}_{c}{\bm{h}}_{c}^{\top},$ (6)

where ${\bm{w}}_{c}$ and ${\bm{h}}_{c}$ denotes column $c$ in ${\bm{W}}$ and
${\bm{H}}$, respectively. The concatenation graph corresponding to an outer
product has a unique path between each pair of row and column nodes. Hence,
performing matrix multiplication by summing over outer products can be viewed
as creating $r$ such concatenation graphs, computing path weights in each
graph separately, and then summing across them.

To convert this into a format suitable for a graph neural network, we need to
accomplish two additional things. First, we need to turn the paths into edges
or nodes. Second, we need to encode the matrix ${\bm{V}}$ that we wish to
factorize.

The first objective can be achieved by converting each outer product graph
${\mathcal{G}}\left({\bm{w}}_{c}{\bm{h}}_{c}^{\top}\right)$ to its
corresponding line digraph
${\mathcal{L}}\left({\mathcal{G}}\left({\bm{w}}_{c}{\bm{h}}_{c}^{\top}\right)\right)$
(Harary & Norman, 1960). This is again a directed bipartite graph with $m$ row
nodes and $n$ column nodes, but the node features in each set correspond to
the components of ${\bm{w}}_{c}$ and ${\bm{h}}_{c}$, respectively, and the
edges correspond to the paths in
${\mathcal{G}}\left({\bm{w}}_{c}{\bm{h}}_{c}^{\top}\right)$.

A directed line graph is unweighted but, since it transforms paths into edges,
it essentially transforms path weights into edge properties. An element in the
outer product can be computed simply by multiplying the node features of the
source and the target nodes. This suggests that if we concatenate the node
features of the line digraphs
${\mathcal{L}}\left({\mathcal{G}}\left({\bm{w}}_{c}{\bm{h}}_{c}^{\top}\right)\right)$
corresponding to each outer product $c=1,\ldots,r$, then an entry $(i,j)$ of
the matrix product ${\bm{W}}{\bm{H}}^{\top}$ can be computed on the edge
${\bm{e}}_{i,j}$ of the resulting graph. Namely, by first computing the
elementwise product of the node features and then summing the result. Further,
if we define ${\bm{e}}_{i,j}={\bm{V}}_{i,j}$, we can compare the current value
of the product with the desired one, thus resolving the second issue of
encoding the matrix ${\bm{V}}$ as well. We will refer to the resulting graph
as an augmented line digraph.

##  5 Optimization Method

There is an inherent symmetry between the factors in the NMF problem since the
objective function (loss) $\ell({\bm{W}},{\bm{H}};{\bm{V}})$ of the
optimization problem in Equation 2 is invariant under transposition

$\ell({\bm{W}},{\bm{H}};{\bm{V}})\triangleq\frac{1}{2}\|{\bm{W}}{\bm{H}}^{\top}-{\bm{V}}\|_{F}^{2}=\ell({\bm{H}},{\bm{W}};{\bm{V}}^{\top}).$
(7)

Again, keeping either ${\bm{W}}$ or ${\bm{H}}$ fixed results in a nonnegative
least-squares problem. Due to symmetry, we can use the same implementation for
both subproblems, and we will therefore only describe the ${\bm{H}}$-update,

$\hat{{\bm{H}}}=\underset{{\bm{H}}\geq
0}{\text{argmin}}\frac{1}{2}\|{\bm{W}}{\bm{H}}^{\top}-{\bm{V}}\|_{\text{F}}^{2}.$
(8)

This can also be interpreted using our graph representation. The
${\bm{H}}$-update corresponds to Figure 2, where the information flows from
${\bm{W}}$ to ${\bm{H}}$. Conversely, the information flows in the opposite
direction in the ${\bm{W}}$-update, which uses ${\bm{V}}^{\top}$.

Following Huang et al. (2016), we apply the alternating direction method of
multipliers (ADMM) (Boyd et al., 2011) to the subproblems, but only unroll a
few iterations of it instead of solving the subproblems exactly at each
iteration. See Appendix B for details. Empirically, such early stopping has
been found to produce better local minima (Udell et al., 2016). In addition to
being faster, other major advantages of adopting ADMM with early stopping for
neural acceleration, as compared to a generic convex optimization layer
(Agrawal et al., 2019), is that it supports warm starting and problems of
variable size.

Various extensions to (i) losses other than the Frobenius norm, and (ii)
constraints other than nonnegativity, are possible within essentially the same
framework (Huang et al., 2016; Udell et al., 2016). We foresee that the neural
acceleration method we outline in the next section would be relatively
straightforward to extend in a similar fashion.

##  6 Neural Acceleration

We propose a neural acceleration method for NMF based on the bipartite graph
representation in Section 4. The overall structure of the model is shown in
Figure 4(a). The acceleration is achieved by a graph neural network that
adapts the exceedingly successful Transformer architecture (Vaswani et al.,
2017) to weighted graphs.

((a)) General model.

((b)) Learned initialization model.

((c)) Learned acceleration model.

Figure 4: We discuss the general model for neural acceleration shown in (a),
and perform experiments using the learned initilization model in (b) and the
neural acceleration model in (c).

Bipartite graphs appear in many application areas, so it is not surprising
that Transformers has previously been extended to this setting (Shu et al.,
2020; Hudson & Zitnick, 2021). However, these works, as well as the original
graph attention network (Veličković et al., 2017), place little (if any)
emphasis on the edge features. In our graph representation, the edge features
are undoubtedly important, since they correspond to the data matrix that we
wish to factorize. To place node and edge features on equal footing, we,
therefore, extend the relation-aware self-attention mechanism of Shaw et al.
(2018) to arbitrary directed, weighted, graphs. We also show how this can be
combined with implicit edge features that are computed dynamically, in what we
refer to as a factor Transformer layer or Factormer for short.

Multiple Factormers are then combined into an $N$-Factormer, which updates the
factors in an alternating fashion using both Factormers and unrolled ADMM
iterations. The constraints are enforced in the final layer by performing a
Euclidean projection on the feasible set, which in the case of NMF simply
amounts to setting negative entries to zero using a rectified linear unit.
More generally, as long as the feasible set is convex, the Euclidean
projection can be expressed as a convex quadratic program that could be
incorporated as a layer in the network (Agrawal et al., 2019).

###  6.1 Implicit Edge Features

According to the graph representation in Section 4, the edge features
correspond to elements in the data matrix ${\bm{e}}_{i,j}={\bm{V}}_{i,j}$, and
are therefore static. However, during execution, we dynamically compute
implicit edge features by augmenting the edge features with the element-wise
products of the node features

$\tilde{{\bm{e}}}_{i,j}({\bm{x}}_{i},{\bm{x}}_{j},{\bm{e}}_{i,j})=\left[{\bm{x}}_{i}\odot{\bm{x}}_{j},{\bm{e}}_{i,j}\right],$
(9)

where $\odot$ denotes element-wise product. This is motivated by the fact that
the residual ${\bm{R}}_{i,j}$ associated with edge $(i,j)$ is given by

$\displaystyle{\bm{R}}_{i,j}$
$\displaystyle=\left(\sum_{l=1}^{r}{\bm{W}}_{i,l}{\bm{H}}_{j,l}\right)-{\bm{V}}_{i,j}$
(10)
$\displaystyle=\tilde{{\bm{e}}}_{i,j}\left({\bm{W}}_{i,:},{\bm{H}}_{j,:}\,,{\bm{V}}_{i,j}\right)\begin{pmatrix}1,\ldots,1,-1\end{pmatrix}^{\top}\\!.$

###  6.2 Factor Transformer Layer

A Factor transformer layer (Factormer) is a message-passing layer that uses
relation-aware multi-head attention with implicit edge features. The input is
a graph with source node features ${\bm{x}}_{i}\in\mathbb{R}^{d}$, target node
features $x_{j}\in\mathbb{R}^{d}$, and edge features
${\bm{e}}_{i,j}\in\mathbb{R}$. In an ${\bm{H}}$-update based on the augmented
line digraph from Section 4, these correspond, in turn, to rows in ${\bm{W}}$,
rows in ${\bm{H}}$, and elements in ${\bm{V}}$. We denote this
${\bm{H}}^{k+1}=\text{Factormer}({\bm{W}}^{k},{\bm{H}}^{k},{\bm{V}})$, where
$k$ is an iteration counter. For clarity of presentation, we omit biases and
indices corresponding to the different heads.

First, we compute key, query, and value as in standard self-attention (Vaswani
et al., 2017)

${\bm{q}}_{j}={\bm{x}}_{j}{\bm{\Theta}}_{\text{Q}},\quad{\bm{k}}_{i}^{\text{N}}={\bm{x}}_{i}{\bm{\Theta}}_{\text{K}}^{\text{N}},\quad{\bm{v}}_{i}^{\text{N}}={\bm{x}}_{i}{\bm{\Theta}}_{V}^{\text{N}},$
(11)

where the weight matrices
${\bm{\Theta}}_{\text{Q}},{\bm{\Theta}}_{\text{K}}^{\text{N}},{\bm{\Theta}}_{\text{V}}^{\text{N}}$
are all square ($d\times d$). We will refer to the usual key and value as node
key and node value, respectively, and indicate that using a superscript N.
This is to distinguish those from the edge key and edge value, with
superscript E, that we compute based on the implicit edge features defined in
equation 9

${\bm{k}}_{i,j}^{\text{E}}=\tilde{{\bm{e}}}_{i,j}{\bm{\Theta}}_{\text{K}}^{\text{E}},\quad{\bm{v}}_{i,j}^{\text{E}}=\tilde{{\bm{e}}}_{i,j}{\bm{\Theta}}_{V}^{\text{E}},$
(12)

where, importantly, the weight matrices
${\bm{\Theta}}_{\text{K}}^{\text{E}},{\bm{\Theta}}_{\text{V}}^{\text{E}}$ map
to the same $d$-dimensional space as the node key and node value.

We compute the message ${\bm{m}}_{j}$ by adding the node and edge values,
multiplying with the attention weight, and finally performing sum aggregation

$\displaystyle\alpha_{i,j}$ $\displaystyle=\text{softmax
}\left({\bm{q}}_{j}\left(\mathbf{k}_{i}^{\text{N}}+\mathbf{k}_{i,j}^{\text{E}}\right)^{\top}\middle/\sqrt{d}\right),$
(13) $\displaystyle{\bm{m}}_{j}$
$\displaystyle=\sum_{i\in\mathcal{N}(j)}\alpha_{i,j}\left({\bm{v}}_{i}^{\text{N}}+{\bm{v}}_{i,j}^{\text{E}}\right).$
(14)

The update function consists of adding the message to the node feature,
applying layer norm LN (Ba et al., 2016), passing it through a fully-connected
feedforward network FFN, and finally adding it to the node feature and
applying layer norm again, that is

${\bm{u}}({\bm{x}}_{j})=\text{LN}\left({\bm{x}}_{j}+\text{FFN}\left(\text{LN}\left({\bm{x}}_{j}+{\bm{m}}_{j}\right)\right)\right).$
(15)

###  6.3 $N$-Factormer

In the spirit of alternating minimization, we arrange multiple Factormers into
a larger unit called an $N$-Factormer by first embedding the rows of
${\bm{W}}$ and ${\bm{H}}$ in a $d$-dimensional space, then composing $2N$
Factormer layers that alternate between updating each of the embedded factors
while maintaining the same dimensionality, and finally mapping back to the
original space

$\displaystyle\tilde{{\bm{H}}}^{0}$
$\displaystyle={\bm{H}}^{\top}{\bm{\Theta}}_{\text{embed}},\quad\tilde{{\bm{W}}}^{0}={\bm{W}}^{\top}{\bm{\Theta}}_{\text{embed}},$
(16) $\displaystyle\tilde{{\bm{H}}}^{1}$
$\displaystyle=\text{Factormer}\left(\tilde{{\bm{W}}}^{0},\tilde{{\bm{H}}}^{0},{\bm{V}}\right),$
$\displaystyle\tilde{{\bm{W}}}^{1}$
$\displaystyle=\text{Factormer}\left(\tilde{{\bm{H}}}^{1},\tilde{{\bm{W}}}^{0},{\bm{V}}^{\top}\right),$
$\displaystyle\vdots$ $\displaystyle\tilde{{\bm{H}}}^{N}$
$\displaystyle=\text{Factormer}\left(\tilde{{\bm{W}}}^{N-1},\tilde{{\bm{H}}}^{N-1},{\bm{V}}\right),$
$\displaystyle\tilde{{\bm{W}}}^{N}$
$\displaystyle=\text{Factormer}\left(\tilde{{\bm{H}}}^{N},\tilde{{\bm{W}}}^{N-1},{\bm{V}}^{\top}\right),$
$\displaystyle{\bm{H}}^{t+1}$
$\displaystyle=\tilde{{\bm{H}}}^{N}{\bm{\Theta}}_{\text{extract}},\quad{\bm{W}}^{t+1}=\tilde{{\bm{W}}}^{N}{\bm{\Theta}}_{\text{extract}},$

where ${\bm{\Theta}}_{\text{embed}}\in\mathbb{R}^{r\times
d},{\bm{\Theta}}_{\text{extract}}\in\mathbb{R}^{d\times r}$ are weight
matrices.

###  6.4 Interleaved Optimization

We evaluate two different acceleration models, “learned intialization” and
“learned acceleration”. Learned initialization only refines the initialization
used to warm start ADMM, while learned acceleration is applied in between the
ADMM updates. We summarize the key ideas below, further details—including
pseudocode—can be found in the Appendix A.

The learned initialization model, shown in Figure 4(b), maps the input
$({\bm{W}}^{0},{\bm{H}}^{0},{\bm{V}})$ through a $N$-Factormer and projects
the output on the nonnegative orthant using a Rectified Linear Unit (ReLU).
The outputs $({\bm{W}}^{1},{\bm{H}}^{1})$ are then further improved using
ADMM.

The learned acceleration model, shown in Figure 4(c), also uses a single
$N$-Factormer and a ReLU, but arranged to perform neural fixed-point
acceleration (Venkataraman & Amos, 2021) for $T$ iterations. At each
iteration, the inputs to the accelerator are (i) the previous output from the
accelerator $({\bm{W}}^{t},{\bm{H}}^{t})$, (ii) the previous output from ADMM
$\left(\hat{{\bm{W}}}^{t},\hat{{\bm{H}}}^{t}\right)$, and (iii) the edge
feature matrix ${\bm{V}}$. Venkataraman & Amos (2021) suggested using a
recurrent neural network for the acceleration, but we did not find any
improved performance from that.

###  6.5 Loss

Without additional assumptions, the NMF solution is not unique even after
removing trivial symmetries (Huang et al., 2013). This makes supervised
learning difficult. Instead we train in an unsupervised fashion, using a
weighted sum of the objective function $\ell({\bm{W}},{\bm{H}};{\bm{V}})$ over
iterations $t=1,\ldots,T$ as the loss (Andrychowicz et al., 2016)

$L\left(\\{{\bm{W}}^{t},{\bm{H}}^{t}\\}_{t=1}^{T},{\bm{V}}\right)=\sum_{k=0}^{T}\gamma^{k}\,\ell({\bm{W}}^{T-k},{\bm{H}}^{T-k},{\bm{V}}),$
(17)

where $\gamma\in(0,1]$ is a discount factor that places more weight on the
loss at later iterations. Note that the final Euclidean projection step of our
network guarantees that the output is always feasible.

##  7 Experiments

We implemented our models using Pytorch Geometric (Fey & Lenssen, 2019). The
graph neural network is initialized using nonnegative SVD (Boutsidis &
Gallopoulos, 2008), as implemented in NIMFA (Žitnik & Zupan, 2012). We applied
the acceleration for $T=5$ iterations and performed 5 ADMM iterations in each
call. Both the learned initialization and the learned acceleration models used
4-Factormers with hidden dimension 100 and 4 self-attention heads. See
Appendix A for further details.

We will release code to reproduce the experiments if the paper is accepted.

###  7.1 Data

####  7.1.1 Synthetic Data

Figure 5: Histograms of pixel intensities for the Yale Face database and the
Jasper Ridge dataset, together with the theoretical probability distribution
of elements in the synthetic matrices.

We train the model on synthetic data matrices ${\bm{V}}$ defined as
${\bm{V}}={\bm{W}}{\bm{H}}^{\top}+{\bm{N}}$, where the elements of ${\bm{W}}$
and ${\bm{H}}$ are sampled i.i.d. according to some predefined distribution
and the noise is i.i.d. normally distributed,
${\textnormal{N}}_{i,j}\sim\mathcal{N}(0,\sigma^{2})$. Obviously, we want the
trained model to generalize to real-world data. Our initial (failed)
experiments, using a Uniform distribution for the elements in ${\bm{W}}$ and
${\bm{H}}$, clearly showed that successful generalization requires a
reasonable similarity between the values of the synthetic and real data
matrices. We thus switched to sampling from an Exponential distribution with
rate parameter $\lambda=1/\sqrt{r}$, chosen so to make the elements in the
data matrix have mean 1,

$\mathbb{E}\left[{\textnormal{V}}_{i,j}\right]=\mathbb{E}\left[{\textnormal{W}}_{i,:}{\textnormal{H}}_{j,:}^{\top}+{\textnormal{N}}_{i,j}\right]=r\cdot\lambda^{2}+0=1.$
(18)

This procedure is similar to the one of Huang et al. (2016) and, as shown in
Figure 5, it roughly matches the distributions of the real data (normalized to
have mean 1).

We created a synthetic training dataset containing 15 000 matrices, out of
which 10 000 were small, $n,m\sim\mathcal{U}[10,35]$, and the remaining 5 000
were larger, $n,m\sim\mathcal{U}[10,100]$ (see Appendix C). The underlying
factors had rank $r=10$ but we also added noise with standard deviation
$\sigma=0.01$. Additional details are provided in Appendix C. Similarly, we
created two noise-free validation datasets: one with 128 small matrices,
$n,m\sim\mathcal{U}[15,35]$, and one with 128 larger matrices,
$n,m\sim\mathcal{U}[50,200]$.

####  7.1.2 Yale Face Database

The Yale Face Database (Belhumeur et al., 1997) contains 165 grayscale images
of size $64\times 64$ pixels depicting 15 individuals. There are 11 images per
subject, one per different facial expression or configuration (center-light,
left-light, with glasses, without glasses, etc.).

####  7.1.3 Hyperspectral imaging

The Jasper Ridge dataset (Rodarmel & Shan, 2002) is one of the most widely
used datasets of hyperspectral images. We consider a sub-image of $100\times
100$ pixels with 198 channels (Zhu, 2017) (spectral bands), where some bands
have been removed due to dense water vapor and atmospheric effects.

###  7.2 Convergence results

Since the matrices have different sizes, we measure the average, element-wise,
reconstruction error of the data matrices, which is equivalent to the root-
mean-square error (RMSE). The RMSE can be interpreted, roughly, as the
standard deviation of the prediction error. To appreciate the scale, recall
that the data matrices have been normalized to have mean 1.

The results on the small and the large synthetic validation datasets are shown
in Figures 6 and 7, respectively. The two leftmost columns are plots of the
RMSE at each iteration, on standard and semi-logarithmic scale. Each of the
two rightmost plots compares the pairwise (same data matrix) RMSE ratio of an
accelerated method to the one without acceleration. The solid line is the
median (second quartile) and the shaded area corresponds to the results in
between the first and third quartiles. The learned initialization performs
notably well, immediately producing a fit that it takes the method without
acceleration 20 iterations to match.

Similarly, the results on the Yale Face database and the Jasper Ridge dataset
are presented in Figures 8 and 9. We also show the 10 basis vectors produced
by the learned initialization followed by 30 ADMM iterations. On both of these
real datasets, the accelerated methods seem to result in a better local
optimum.

((a))

((b))

((c))

((d))

Figure 6: Results on the small synthetic validation dataset for the two
accelerated methods and the baseline AO-ADMM methodl. a) The reconstruction
error (RMSE) at each iteration; b) $\log_{10}$ RMSE; c): the median and the
quartiles of the pairwise ratio of reconstruction errors between the learned
initialization model and the baseline; d): the median and the quartiles of the
pairwise ratio of reconstruction errors between the learned acceleration model
and the baseline.

((a))

((b))

((c))

((d))

Figure 7: Results on the large synthetic validation dataset for the two
accelerated methods and the baseline AO-ADMM method. a) reconstruction error
(RMSE) at each iteration; b) $\log_{10}$ RMSE; c): the median and the
quartiles of the pairwise ratio of reconstruction errors between the learned
initialization model and the baseline; d): the median and the quartiles of the
pairwise ratio of reconstruction errors between the learned acceleration model
and the baseline.

((a))

((b))

((c))

Figure 8: Yale Face Database. a): reconstruction errors (RMSE) at each
iteration for the two accelerated methods and the baseline AO-ADMM method; b)
: Pairwise ratios of the reconstruction error of each accelerated method to
the baseline; c): images corresponding to the ten basis vectors produced by
the learned initialization model followed by 30 iterations of ADMM.

((a))

((b))

((c))

Figure 9: Jasper ridge hyperspectral data set. a) Convergence plots for the
two accelerated methods and the baseline AO-ADMM method; b) Pairwise ratios of
the reconstruction error of each accelerated method to the baseline; c):
Images corresponding to the ten basis vectors produced by the learned
initialization model followed by 30 iterations of ADMM.

##  8 Limitations

There are two major limitations of this work that we are aware of. The first
is that, since the rank corresponds directly to the dimensionality of the node
features, the network can only be trained for a given rank. Potentially, this
could be addressed by borrowing ideas used to create node embeddings (Hamilton
et al., 2017). The second major limitation is computational scalability.
Transformers in general are known to be computationally expensive, although
remedies have been proposed (Kitaev et al., 2020). But another bottleneck is
that the augmented line digraph is a _complete_ bipartite graph where the
nodes either have degree $m$ or $n$, which makes the message-passing step
computationally costly for large matrices. This is, however, a problem shared
with many real-world graph problems (Li et al., 2021), which has therefore led
to several promising strategies, including neighborhood sampling (Hamilton et
al., 2017; Chen et al., 2018) and graph sampling (Zeng et al., 2020).

##  9 Conclusions

In this work, we have shown how the König digraph gives us a way of designing
graph neural networks for numerical linear algebra and exemplified that by
proposing an augmented line digraph suited to low-rank factorization problems.
We further described a graph neural network architecture that interleaves
Factormers with unrolled optimization steps and showed that it could
substantially accelerate alternating optimization schemes for nonnegative
matrix factorization. We foresee that recent and ongoing research on graph
neural networks will lead to even more capable architectures, as well as
clever ideas to improve computational scalability, which would further extend
the applicability of this work.

## References

  * Agrawal et al. (2019) Agrawal, A., Amos, B., Barratt, S., Boyd, S., Diamond, S., and Kolter, J. Z.  Differentiable convex optimization layers.  _Advances in Neural Information Processing Systems_ , 32:9562–9574, 2019.
  * Albright et al. (2006) Albright, R., Cox, J., Duling, D., Langville, A. N., and Meyer, C.  Algorithms, initializations, and convergence for the nonnegative matrix factorization.  Technical report, Tech. rep. 919. NCSU Technical Report Math 81706. http://meyer. math. ncsu …, 2006.
  * Andrychowicz et al. (2016) Andrychowicz, M., Denil, M., Colmenarejo, S. G., Hoffman, M. W., Pfau, D., Schaul, T., Shillingford, B., and de Freitas, N.  Learning to learn by gradient descent by gradient descent.  In _Proceedings of the 30th International Conference on Neural Information Processing Systems_ , pp. 3988–3996, 2016.
  * Ba et al. (2016) Ba, J. L., Kiros, J. R., and Hinton, G. E.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Banert et al. (2020) Banert, S., Ringh, A., Adler, J., Karlsson, J., and Oktem, O.  Data-driven nonsmooth optimization.  _SIAM Journal on Optimization_ , 30(1):102–131, 2020.
  * Belhumeur et al. (1997) Belhumeur, P. N., Hespanha, J. P., and Kriegman, D. J.  Eigenfaces vs. fisherfaces: Recognition using class specific linear projection.  _IEEE Transactions on pattern analysis and machine intelligence_ , 19(7):711–720, 1997.
  * Bello et al. (2017) Bello, I., Zoph, B., Vasudevan, V., and Le, Q. V.  Neural optimizer search with reinforcement learning.  In _International Conference on Machine Learning_ , pp. 459–468. PMLR, 2017.
  * Berg et al. (2017) Berg, R. v. d., Kipf, T. N., and Welling, M.  Graph convolutional matrix completion.  _arXiv preprint arXiv:1706.02263_ , 2017.
  * Boutsidis & Gallopoulos (2008) Boutsidis, C. and Gallopoulos, E.  SVD based initialization: A head start for nonnegative matrix factorization.  _Pattern recognition_ , 41(4):1350–1362, 2008\.
  * Boyd et al. (2011) Boyd, S., Parikh, N., Chu, E., Peleato, B., and Eckstein, J.  Distributed optimization and statistical learning via the alternating direction method of multipliers.  _Foundations and Trends® in Machine Learning_ , 3(1):1–122, 2011.
  * Bronstein et al. (2017) Bronstein, M. M., Bruna, J., LeCun, Y., Szlam, A., and Vandergheynst, P.  Geometric deep learning: going beyond euclidean data.  _IEEE Signal Processing Magazine_ , 34(4):18–42, 2017.
  * Brualdi & Cvetkovic (2008) Brualdi, R. A. and Cvetkovic, D.  _A combinatorial approach to matrix theory and its applications_.  CRC press, 2008.
  * Cappart et al. (2021) Cappart, Q., Chételat, D., Khalil, E., Lodi, A., Morris, C., and Veličković, P.  Combinatorial optimization and reasoning with graph neural networks.  _arXiv preprint arXiv:2102.09544_ , 2021.
  * Chen et al. (2018) Chen, J., Zhu, J., and Song, L.  Stochastic training of graph convolutional networks with variance reduction.  In _International Conference on Machine Learning_ , pp. 942–950. PMLR, 2018.
  * Chen et al. (2021) Chen, T., Chen, X., Chen, W., Heaton, H., Liu, J., Wang, Z., and Yin, W.  Learning to optimize: A primer and a benchmark.  _arXiv preprint arXiv:2103.12828_ , 2021.
  * Dai et al. (2017) Dai, H., Khalil, E., Zhang, Y., Dilkina, B., and Song, L.  Learning combinatorial optimization algorithms over graphs.  _Advances in Neural Information Processing Systems_ , 30:6348–6358, 2017.
  * Diamond et al. (2017) Diamond, S., Sitzmann, V., Heide, F., and Wetzstein, G.  Unrolled optimization with deep priors.  _arXiv preprint arXiv:1705.08041_ , 2017.
  * Domke (2012) Domke, J.  Generic methods for optimization-based modeling.  In _Artificial Intelligence and Statistics_ , pp. 318–326. PMLR, 2012.
  * Doob (1984) Doob, M.  Applications of graph theory in linear algebra.  _Mathematics Magazine_ , 57(2):67–76, 1984.
  * Duff et al. (2017) Duff, I. S., Erisman, A. M., and Reid, J. K.  _Direct methods for sparse matrices_.  Oxford University Press, 2017.
  * Fey & Lenssen (2019) Fey, M. and Lenssen, J. E.  Fast graph representation learning with PyTorch Geometric.  _arXiv preprint arXiv:1903.02428_ , 2019.
  * Fu et al. (2019) Fu, X., Huang, K., Sidiropoulos, N. D., and Ma, W.-K.  Nonnegative matrix factorization for signal and data analytics: Identifiability, algorithms, and applications.  _IEEE Signal Process. Mag._ , 36(2):59–80, 2019\.
  * Gasse et al. (2019) Gasse, M., Chételat, D., Ferroni, N., Charlin, L., and Lodi, A.  Exact combinatorial optimization with graph convolutional neural networks.  _Advances in Neural Information Processing Systems_ , 33, 2019.
  * Gillis (2020) Gillis, N.  _Nonnegative Matrix Factorization_.  SIAM, 2020.  ISBN 978-1-61197-640-3.
  * Gillis & Vavasis (2013) Gillis, N. and Vavasis, S. A.  Fast and robust recursive algorithms for separable nonnegative matrix factorization.  _IEEE transactions on pattern analysis and machine intelligence_ , 36(4):698–714, 2013.
  * Gilmer et al. (2017) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In _Proceedings of the 34th International Conference on Machine Learning (ICML)_ , 2017.
  * Gregor & LeCun (2010) Gregor, K. and LeCun, Y.  Learning fast approximations of sparse coding.  In _Proceedings of the 27th International Conference on Machine Learning_ , pp. 399–406, 2010.
  * Hamilton et al. (2017) Hamilton, W. L., Ying, R., and Leskovec, J.  Inductive representation learning on large graphs.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , pp. 1025–1035, 2017.
  * Harary & Norman (1960) Harary, F. and Norman, R. Z.  Some properties of line digraphs.  _Rendiconti del Circolo Matematico di Palermo_ , 9(2):161–168, 1960.
  * Huang et al. (2013) Huang, K., Sidiropoulos, N. D., and Swami, A.  Non-negative matrix factorization revisited: Uniqueness and algorithm for symmetric decomposition.  _IEEE Transactions on Signal Processing_ , 62(1):211–224, 2013.
  * Huang et al. (2016) Huang, K., Sidiropoulos, N. D., and Liavas, A. P.  A flexible and efficient algorithmic framework for constrained matrix and tensor factorization.  _IEEE Transactions on Signal Processing_ , 64(19):5052–5065, 2016.
  * Hudson & Zitnick (2021) Hudson, D. A. and Zitnick, C. L.  Generative adversarial transformers.  _Proceedings of the 38th International Conference on Machine Learning, ICML 2021_ , 2021.
  * Hutter et al. (2019) Hutter, F., Kotthoff, L., and Vanschoren, J.  _Automated machine learning: methods, systems, challenges_.  Springer Nature, 2019.
  * Kepner & Gilbert (2011) Kepner, J. and Gilbert, J.  _Graph algorithms in the language of linear algebra_.  SIAM, 2011.
  * Kim & Park (2007) Kim, H. and Park, H.  Sparse non-negative matrix factorizations via alternating non-negativity-constrained least squares for microarray data analysis.  _Bioinformatics_ , 23(12):1495–1502, 2007.
  * Kim & Park (2008) Kim, H. and Park, H.  Nonnegative matrix factorization based on alternating nonnegativity constrained least squares and active set method.  _SIAM journal on matrix analysis and applications_ , 30(2):713–730, 2008.
  * Kipf & Welling (2017) Kipf, T. N. and Welling, M.  Semi-supervised classification with graph convolutional networks.  In _5th International Conference on Learning Representations (ICLR)_ , 2017.
  * Kitaev et al. (2020) Kitaev, N., Kaiser, L., and Levskaya, A.  Reformer: The efficient transformer.  In _International Conference on Learning Representations_ , 2020.
  * Lee & Seung (1999) Lee, D. D. and Seung, H. S.  Learning the parts of objects by non-negative matrix factorization.  _Nature_ , 401(6755):788–791, 1999.
  * Lee & Seung (2001) Lee, D. D. and Seung, H. S.  Algorithms for non-negative matrix factorization.  In _Advances in neural information processing systems_ , pp. 556–562, 2001.
  * Li et al. (2021) Li, G., Müller, M., Ghanem, B., and Koltun, V.  Training graph neural networks with 1000 layers.  In _Proceedings of the 38th International Conference on Machine Learning_ , volume 139. PMLR, 2021.
  * Li & Malik (2016) Li, K. and Malik, J.  Learning to optimize.  _arXiv preprint arXiv:1606.01885_ , 2016.
  * Lin (2007) Lin, C.-J.  Projected gradient methods for nonnegative matrix factorization.  _Neural computation_ , 19(10):2756–2779, 2007\.
  * Loshchilov & Hutter (2016) Loshchilov, I. and Hutter, F.  Sgdr: Stochastic gradient descent with warm restarts.  _arXiv preprint arXiv:1608.03983_ , 2016.
  * Loshchilov & Hutter (2017) Loshchilov, I. and Hutter, F.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Maheswaranathan et al. (2020) Maheswaranathan, N., Sussillo, D., Metz, L., Sun, R., and Sohl-Dickstein, J.  Reverse engineering learned optimizers reveals known and novel mechanisms.  _arXiv preprint arXiv:2011.02159_ , 2020.
  * Monga et al. (2021) Monga, V., Li, Y., and Eldar, Y. C.  Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing.  _IEEE Signal Processing Magazine_ , 38(2):18–44, 2021.
  * Monti et al. (2017) Monti, F., Bronstein, M. M., and Bresson, X.  Geometric matrix completion with recurrent multi-graph neural networks.  _Proceedings of the 34th International Conference on Machine Learning (ICML)_ , 2017.
  * Nair et al. (2020) Nair, V., Bartunov, S., Gimeno, F., von Glehn, I., Lichocki, P., Lobov, I., O’Donoghue, B., Sonnerat, N., Tjandraatmadja, C., Wang, P., et al.  Solving mixed integer programs using neural networks.  _arXiv preprint arXiv:2012.13349_ , 2020.
  * Recht et al. (2012) Recht, B., Re, C., Tropp, J., and Bittorf, V.  Factoring nonnegative matrices with linear programs.  _Advances in neural information processing systems_ , 25:1214–1222, 2012.
  * Rodarmel & Shan (2002) Rodarmel, C. and Shan, J.  Principal component analysis for hyperspectral image classification.  _Surveying and Land Information Science_ , 62(2):115–122, 2002.
  * Shahnaz et al. (2006) Shahnaz, F., Berry, M. W., Pauca, V. P., and Plemmons, R. J.  Document clustering using nonnegative matrix factorization.  _Information Processing & Management_, 42(2):373–386, 2006.
  * Shaw et al. (2018) Shaw, P., Uszkoreit, J., and Vaswani, A.  Self-attention with relative position representations.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers)_ , pp. 464–468, 2018.
  * Shu et al. (2020) Shu, X., Xue, M., Li, Y., Zhang, Z., and Liu, T.  BiG-transformer: Integrating hierarchical features for transformer via bipartite graph.  In _2020 International Joint Conference on Neural Networks (IJCNN)_ , pp. 1–8. IEEE, 2020.
  * Udell et al. (2016) Udell, M., Horn, C., Zadeh, R., Boyd, S., et al.  Generalized low rank models.  _Foundations and Trends® in Machine Learning_ , 9(1):1–118, 2016.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  In _Advances in neural information processing systems_ , pp. 5998–6008, 2017.
  * Vavasis (2010) Vavasis, S. A.  On the complexity of nonnegative matrix factorization.  _SIAM Journal on Optimization_ , 20(3):1364–1377, 2010.
  * Veličković et al. (2017) Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y.  Graph attention networks.  _arXiv preprint arXiv:1710.10903_ , 2017.
  * Venkataraman & Amos (2021) Venkataraman, S. and Amos, B.  Neural fixed-point acceleration for convex optimization.  _8th ICML Workshop on Automated Machine Learning_ , 2021.
  * Xu & Yin (2013) Xu, Y. and Yin, W.  A block coordinate descent method for regularized multiconvex optimization with applications to nonnegative tensor factorization and completion.  _SIAM Journal on imaging sciences_ , 6(3):1758–1789, 2013.
  * Zeng et al. (2020) Zeng, H., Zhou, H., Srivastava, A., Kannan, R., and Prasanna, V.  GraphSAINT: Graph sampling based inductive learning method.  In _International Conference on Learning Representations_ , 2020.
  * Zhu (2017) Zhu, F.  Hyperspectral unmixing: Ground truth labeling, datasets, benchmark performances and survey.  _arXiv: Computer Vision and Pattern Recognition_ , 2017.
  * Žitnik & Zupan (2012) Žitnik, M. and Zupan, B.  NIMFA: A Python library for nonnegative matrix factorization.  _The Journal of Machine Learning Research_ , 13(1):849–853, 2012.

##  Appendix A Implementation details

We present the pseudo-code of the learned initialization model and the learned
acceleration model in Algorithm 1 and 2. We present the pseudo-code for the
Factormer in Algorithm 3.

We used the AdamW optimizer (Loshchilov & Hutter, 2017) for all experiments.
We set the batch size to one. We trained the models using cosine annealing
warm start (Loshchilov & Hutter, 2016) where we let the learning rate decrease
during the whole epoch. We set the initial learning rate to $10^{-4}$ for the
initialization model and $10^{-5}$ for the learned acceleration model. After
every epoch, we decreased the initial learning rate by a factor of 0.9. We
trained both models during 15 epochs. For the learned acceleration model, we
set the discount factor on the loss to $\gamma=0.2$.

When we trained the acceleration model, we first used one acceleration step.
We then increased the number of acceleration steps by one at every two epochs
until we reached five acceleration iterations.

Algorithm 1 Learned initialization model

Input: $\mathbf{W}^{0}$, $\mathbf{H}^{0}$, ${\bm{V}}$, $T$

${\bm{H}}^{0}=\text{LinearEmbed}({\bm{H}}^{0})$,
${\bm{W}}^{0}=\text{LinearEmbed}({\bm{W}}^{0})$ {maps rank to hidden
dimension.}

for Factormer in N-Factormer do

$\mathbf{H}^{0}=\text{Factormer}({\bm{W}}^{0},{\bm{H}}^{0},\mathbf{V})$

$\mathbf{W}^{0}=\text{Factormer}(\mathbf{H}^{0},\mathbf{W}^{0},\mathbf{V}^{T})$

end for

$\mathbf{H}^{1}=\text{LinearExtract}(\mathbf{H}^{0})$ , $\
\mathbf{W}^{1}=\text{LinearExtract}(\mathbf{W}^{0})$ {maps hidden dimension to
rank.}

for $t=1\ldots T$ do

${{\bm{H}}}^{t+1}=\text{Solver}({\bm{W}}^{t},{\bm{H}}^{t},{\bm{V}})$

${{\bm{W}}}^{t+1}=\text{Solver}({\bm{H}}^{t+1},{\bm{W}}^{t},{\bm{V}}^{T})$

end for

Algorithm 2 Learned acceleration model

Input: $\mathbf{W}^{0}$, $\mathbf{H}^{0}$, ${\bm{V}}$, $T$, nbrAcc {nbrAcc =
number of acceleration steps}

for $t=0\ldots T-1$ do

$\hat{{\bm{H}}}^{t}=\text{Solver}({\bm{W}}^{t},{\bm{H}}^{t},{\bm{V}})$

$\hat{{\bm{W}}}^{t}=\text{Solver}(\hat{{\bm{H}}}^{t},{\bm{W}}^{t},{\bm{V}}^{T})$

if $t<\text{nbrAcc}$ then

$\mathbf{H}^{t}=\text{Concatenate}(\mathbf{H}^{t},\widehat{\mathbf{H}}^{t})$,
$\ \mathbf{W}^{t}=\text{Concatenate}(\mathbf{W}^{t},\hat{\mathbf{W}}^{t})$

$\mathbf{H}^{t}=\text{LinearEmbed}(\mathbf{H}^{t})$ , $\
\mathbf{W}^{t}=\text{LinearEmbed}(\mathbf{W}^{t})$ {maps rank$\times 2$ to
hidden dimension.}

for Factormer in N-Factormer do

$\mathbf{H}^{t}=\text{Factormer}(\mathbf{W}^{t},\mathbf{H}^{t},\mathbf{V})$

$\mathbf{W}^{t}=\text{Factormer}(\mathbf{H}^{t},\mathbf{W}^{t},\mathbf{V}^{T})$

end for

$\mathbf{H}^{t+1}=\text{LinearExtract}(\mathbf{H}^{t})$, $\
\mathbf{W}^{t+1}=\text{LinearExtract}(\mathbf{W}^{t})$ {maps hidden dimension
to rank.}

else

$\mathbf{H}^{t+1}=\hat{{\bm{H}}}^{t}$

$\mathbf{W}^{t+1}=\hat{{\bm{W}}}^{t}$

end if

end for

Algorithm 3 Factormer

Input: $\mathbf{x}_{i}$, $\mathbf{x}_{j}$, $\mathbf{e}_{i,j}$ {Source nodes,
target nodes and edge attributes.}

Derive message:

$\tilde{\mathbf{e}}_{i,j}=\left[{\bm{x}}_{i}\odot{\bm{x}}_{j},{\bm{e}}_{i,j}\right]$

${\bm{q}}_{j}=\text{LinearQuery}({\bm{x}}_{j})$,
${\bm{k}}_{i}^{N}=\text{LinearNodeKey}({\bm{x}}_{i})$,
${\bm{k}}_{i}^{E}=\text{LinearEdgeKey}({\bm{e}}_{i,j})$

${\bm{v}}_{i}^{N}=\text{LinearNodeValue}({\bm{x}}_{i})$,
${\bm{v}}_{i}^{E}=\text{LinearEdgeValue}({\bm{e}}_{i,j})$

$\alpha_{i,j}=\text{softmax
}\left({\bm{q}}_{j}\left(\mathbf{k}_{i}^{\text{N}}+\mathbf{k}_{i,j}^{\text{E}}\right)^{\top}\middle/\sqrt{d}\right)$
{Derive self attention.}

$\mathbf{m}_{i,j}=\alpha_{i,j}\left({\bm{v}}_{i}^{\text{N}}+{\bm{v}}_{i,j}^{\text{E}}\right)${Derive
message.}

Aggregate:

${\bm{x}}_{j}^{\prime}=\sum_{i\in\mathcal{N}(j)}{\bm{m}}_{i,j}$ {Add
aggregation.}

Update:

$\mathbf{x}_{j}=\mathbf{x}_{j}^{\prime}+\mathbf{x}_{j}$

$\mathbf{x}_{j}=\text{LayerNorm}(\mathbf{x}_{j})$

$\mathbf{x}_{j}=\text{FeedForward}(\mathbf{x}_{j})$

$\mathbf{x}_{j}=\mathbf{x}_{j}^{\prime}+\mathbf{x}_{j}$

if not last Factormer then

$\mathbf{x}_{j}=\text{LayerNorm}(\mathbf{x}_{j})${In the last N-Factormer
layer we skip the last layer norm.}

end if

##  Appendix B ADMM for nonnegative least-squares

To derive the alternating direction method of multipliers (ADMM) equations for
the nonnegative least-squares problem in Equation 8, we introduce an auxiliary
variable $\tilde{{\bm{H}}}\in\mathbb{R}^{n\times r}$ and rewrite the problem
as

$\displaystyle\underset{{\bm{H}}}{\text{minimize}}$
$\displaystyle\frac{1}{2}\|{\bm{W}}{\bm{H}}^{\top}-{\bm{V}}\|_{\text{F}}^{2}+I_{\tilde{{\bm{H}}}\geq
0}(\tilde{{\bm{H}}})$ (19) subject to
$\displaystyle{\bm{H}}-\tilde{{\bm{H}}}=0,$

where $I_{\tilde{{\bm{H}}}\geq 0}$ is an indicator function which is zero when
$\tilde{{\bm{H}}}\geq 0$ and infinite otherwise. Using an augmented Lagrangian
parameter $\rho>0$ (we use $\rho=1$) and a scaled dual variable
${\bm{U}}\in\mathbb{R}^{n\times r}$, an iteration of ADMM then consists of the
following equations,

$\displaystyle{\bm{H}}$
$\displaystyle\leftarrow\left({\bm{W}}^{\top}{\bm{W}}+\rho
I\right)^{-1}\left({\bm{W}}^{\top}{\bm{V}}+\rho(\tilde{{\bm{H}}}+{\bm{U}})\right),$
(20) $\displaystyle\tilde{{\bm{H}}}$
$\displaystyle\leftarrow\left({\bm{H}}-{\bm{U}}\right)_{+},$ (21)
$\displaystyle{\bm{U}}$
$\displaystyle\leftarrow{\bm{U}}+{\bm{H}}-\tilde{{\bm{H}}},$ (22)

where $(\cdot)_{+}$ denotes the nonnegative part of the expression, i.e. a
projection on the nonnegative orthant. For computational efficiency, the
factorization of ${\bm{W}}^{\top}{\bm{W}}+\rho I$ can be cached.

##  Appendix C Synthetic dataset

We created the training data by concatenating small and large matrices. In
Table 1 we describe in detail which matrices as included in the training data.

Table 1: The sizes of the matrices in the training dataset. The total number
of samples in the training dataset is 15000. The number or rows and columns in
the different matrices are sampled uniformly as
$n,m\sim\mathcal{U}[low,high]$.
Nbr of samples | Row range | Column range
---|---|---
10000 | $[10,35]$ | $(10,35)$
2000 | [30, 70] | [30, 70]
1500 | [30,100] | [10,35]
1500 | [10,35] | [30,100]

##  Appendix D Additional results

The root-mean square error (RMSE) is related to the Frobenius norm as follows

$\text{RMSE}=\|{\bm{W}}{\bm{H}}^{\top}-{\bm{V}}\|_{\text{F}}/\sqrt{mn}.$ (23)

((a)) Only nonnegative SVD.

((b)) Initialization network followed by two ADMM iteration.

((c)) Initialization network followed by 30 ADMM iterations.

Figure 10: Images corresponding to the ten basis vectors produced by: only
nonnegateve SVD (a), the learned initialization model followed by two ADMM
iterations (b), the learned initialization model followed by 30 iterations of
ADMM.