# Boolean Matrix Factorization via Nonnegative Auxiliary Optimization

Duc P. Truong
Computer, Computational and Statistics Division
Los Alamos National Laboratory, USA
<EMAIL_ADDRESS>
&Erik Skau
Computer, Computational and Statistics Division
Los Alamos National Laboratory, USA
email
&Derek Desantis
Theoretical Division
Los Alamos National Laboratory, USA
email
&Boian Alexandrov
Theoretical Division
Los Alamos National Laboratory, USA
email

###### Abstract

A novel approach to Boolean matrix factorization (BMF) is presented. Instead
of solving the BMF problem directly, this approach solves a nonnegative
optimization problem with the constraint over an auxiliary matrix whose
Boolean structure is identical to the initial Boolean data. Then the solution
of the nonnegative auxiliary optimization problem is thresholded to provide a
solution for the BMF problem. We provide the proofs for the equivalencies of
the two solution spaces under the existence of an exact solution. Moreover,
the nonincreasing property of the algorithm is also proven. Experiments on
synthetic and real datasets are conducted to show the effectiveness and
complexity of the algorithm compared to other current methods.

##  1 Introduction

Many research fields, such as personalized medicine, space research, social
sciences, climate research, nonproliferation, emergency response, finances,
etc. collect and generate large-scale datasets. Analysis of such datasets is
difficult, since often the underlying fundamental processes (or features)
remain hidden or latent (i.e., not directly observable) [4]. Extracting such
latent features reveals valuable information and hidden causality and
relations. One of the most powerful tools for extracting latent (hidden)
features from data is factor analysis [13]. Traditionally, factor analysis
approximates a data-matrix $X\in\mathbb{R}^{n\times m}$ by a product of two
low-rank (factor) matrices, $X\approx WH$, where $W\in\mathbb{R}^{n\times k}$,
$H\in\mathbb{R}^{k\times m}$, and $k\ll n,m$. Various types of factorizations
can be obtained by imposing different constraints. For instance, imposing
orthogonality on the factors leads to the singular value decomposition (SVD)
[14], while the nonnegativity constraint leads to non-negative matrix
factorization (NMF) [5]. In a lot of investigations, the variables are simple
dichotomies, {false, true}, that is, the data contains only binary values,
$\\{0,1\\}$. For example, in relational databases, an object-attribute
relation is represented by a Boolean variable, which takes value
$\\{1\\}$-true, if the object has this attribute, or $\\{0\\}$-false,
otherwise. In this case, all the values of the factors have to be $0$ or $1$.
In this case, instead of simple arithmetic we also need to use Boolean
algebra: instead of “plus” and “times” we need the logical operations of “or”
with “and”, which results in a Boolean factorization [8]. To be precise,
letting $\mathbb{B}=\\{0,1\\}$, and assuming the input data is a matrix
$X\in\mathbb{B}^{n\times m}$, where $\mathbb{B}=\\{0,1\\}$, Boolean matrix
factorization (BMF) is looking for two binary matrices
$W\in\mathbb{B}^{n\times k}$, $H\in\mathbb{B}^{k\times m}$ such that $X\approx
W\otimes_{B}H$, where $(W\otimes_{B}H)_{ij}=\vee_{l=1}^{k}W_{il}\land
H_{lj}\in\mathbb{B}$ and $\vee$ and $\land$ are the logical ”or” and ”and”
operations. The Boolean rank of $X$ is the smallest $k$ for which such an
exact representation, $X=W\otimes_{B}H$, exists. Interestingly, in contrast to
the non-negative rank, the Boolean rank can be not only bigger or equal, but
also smaller than the real rank [10, 2].

In this paper, we propose a new algorithm for BMF, we call Boolean Auxiliary
Non-negative Matrix Factorization (BANMF). Here, we consider a specific
constrained optimization over an auxiliary matrix $Y$ related to the initial
binary data matrix $X$. By performing the constrained optimization over the
auxiliary matrix we are able to simulate a Boolean matrix product by enforcing
the support of the matrix product and disregarding the exact matrix values.
From the solution of our relaxed problem we apply thresholding to our solution
matrices $W$ and $H$ to arrive at the BMF solution. We demonstrate that our
method outperforms other state-of-the-art methods for the BMF problem when
there exists a discrepancy between the nonnegative and Boolean rank. Moreover,
BANMF has superior performance under corruption by Bernoulli noise and at
various density levels. Furthermore BANMF is shown to have faster run times
than PNL-PF, a competitor algorithm, for large datasets.

##  2 Related Work

The BMF problem is NP-hard [8], leading to the development of a variety of
strategies to approximate a solution. Greedy algorithms were some of the first
proposed methods to solving BMF [8, 1, 7]. There is also a family of Bayesian
methods. For example, the authors in [12] use a probabilistic generative model
and derive a sampling method to accelerate the optimization for BMF. In [11]
they use a Bayesian framework to consider the BMF problem as a maximum log
likelihood problem and use a message passing procedure to approximate the MAP
assignment.

An alternative approach to BMF is to relax the Boolean constraint into a non-
negativity constraint. Several different methods to relax BMF to a NMF problem
exist. The thresholding method in [15] attempts take an NMF solution to a
binary matrix factorization solution. This idea can easily be modified to
search for a BMF solution by searching for a optimal thresholding values. The
Post Nonlinear Penalty Function Algorithm (PNLPF) [9] attempts to improve the
relaxation by putting a nonlinear function on the product WH. The nonlinear
Heaviside function clips arbitrarily large numbers in the product WH back to
the [0,1] interval. This composition of regular matrix product and the
nonlinear function approximates a Boolean matrix product resulting in a BMF
estimation algorithm.

##  3 Proposed Approach

Given a Boolean matrix $X\in\mathbb{B}^{M\times N}$, the BMF optimization
problem is

$\displaystyle\underset{W,H}{\text{minimize}}$
$\displaystyle||X-W\operatorname{\otimes}_{B}H||_{F}$ (1) subject to
$\displaystyle W\in\mathbb{B}^{M\times k},$ $\displaystyle
H\in\mathbb{B}^{k\times N},$

where $\operatorname{\otimes}_{B}$ denotes Boolean matrix product
$(W\operatorname{\otimes}_{B}H)_{ij}=\lor_{l=1}^{k}W_{il}\land H_{lj}$, and
$||...||_{F}$ is the Frobenius norm.

###  3.1 Nonnegative auxiliary optimization for BMF

Here we propose an alternative nonnegative optimization problem to solve the
BMF problem (1) which is comprised of two steps. The first step is based on
NMF optimization with additional constrains on an auxiliary matrix $Y$. Given
a solution to the first step, the second step consists of thresholding to
obtain a Boolean factors.

We now describe the nonnegative auxiliary optimization problem. Given a binary
matrix $X$, we search for a rank $k$ Boolean decomposition by solving the
optimization,

$\displaystyle\underset{Y,W,H}{\text{minimize}}$ $\displaystyle||Y-WH||_{F}$
(2) subject to $\displaystyle 1\leq Y_{i,j}\leq k,\text{ if }X_{i,j}=1$
$\displaystyle Y_{i,j}=0,\text{ if }X_{i,j}=0$ $\displaystyle
W\in\mathbb{R}_{+}^{N\times k},$ $\displaystyle H\in\mathbb{R}_{+}^{k\times
M}.$

The solution $(Y,W,H)$ of problem (2) is not necessarily Boolean. For the
second step, we convert the non-negative $(W,H)$ to Boolean matrices
$(\hat{W},\hat{H})$ according to a thresholding function. Given threshold
parameter $\delta\geq 0$, we define

$\hat{C}_{ij}=\begin{cases}\begin{aligned} 1&\quad\text{ if }C_{ij}>\delta\\\
0&\quad\text{ if }C_{ij}\leq\delta\end{aligned}\end{cases}$ (3)

In theory, we use $\delta=0$. However as will be discussed in Algorithm 2,
$\delta>0$ in practice.

Our method can be viewed as an NMF approach for BMF problem similar to the
other relaxation approaches in the literature [15, 9]. However, by adding the
auxiliary matrix constraint, the BANMF approach forces the support of NMF to
adhere to that of a BMF while disregarding the exact values. The subsequent
results establish this. Under the existence of an exact decomposition, below
we prove the equivalencies of the two solution spaces. This will be done with
two theorems, one theorem showing containment in each direction. The following
theorem proves that an exact decomposition solution of BMF problem (1) is also
a solution of our relaxed nonnegative optimization problem (2).

######  Theorem 1.

Given an exact decomposition solution $(W,H)$ to problem (1) such that
$X=W\operatorname{\otimes}_{B}H$, there exists $Y\in\mathbb{R}_{+}^{N\times
M}$ such that $(Y,W,H)$ is an exact decomposition solution for the
optimization problem (2).

###### Proof.

Given an an exact solution $(W,H)$ such that $X=W\operatorname{\otimes}_{B}H$,
we construct a $N\times M$ nonnegative matrix $Y=WH$.

Now we need to show that $(Y,W,H)$ is an exact solution for the optimization
problem (2). First, by construction, the optimized value is already zero.
Next, we need to show that the solution satisfies the constraints.
(a) $W\in\mathbb{B}^{N\times k}\Rightarrow W\in\mathbb{R}_{+}^{N\times k}$
(b) $H\in\mathbb{B}^{k\times M}\Rightarrow H\in\mathbb{R}_{+}^{k\times M}$
(c) For $(i,j)\notin\operatorname{support}(X)$, which means $X_{ij}=0$
$\Rightarrow W_{il}=H_{lj}=0\ \forall l=1,\dots,k.$ by the definition of
Boolean matrix product in Equation 1
$\Rightarrow Y_{ij}=(WH)_{ij}=\sum_{l=1}^{k}(W)_{il}(H)_{lj}=0$
(d) For $(i,j)\in\operatorname{support}(X)$, which means $X_{ij}=1$
$\Rightarrow\exists l\in\\{1,\dots,k\\}$ such that $W_{il}=H_{lj}=1$. Since
$W_{i,:},\ H_{:,j}\in\mathbb{R}^{k}$, we have $1\leq(WH)_{ij}\leq k\Rightarrow
1\leq Y_{ij}=(WH)_{ij}\leq k$

Therefore, from (a),(b),(c), (d) and the fact that the optimized value is
zero, we have shown that $Y=WH$ is an exact solution for the optimization
problem (2) as needed. ∎

This proves that the BMF solution set is contained in the BANMF solution set
when there exists an exact decomposition. The next theorem proves containment
the other direction, that the BANMF solution set, after Booleanizing the
solutions, is contained in the BMF solution set when there exists an exact
decomposition.

######  Theorem 2.

Given an exact solution $Y=WH$ for the BANMF problem (2), and $\hat{W}$ and
$\hat{H}$ are the Booleanized matrices of $W$ and $H$ with threshold
$\delta=0$, then $X=\hat{W}\operatorname{\otimes}_{B}\hat{H}$ is an exact
solution for the BMF problem (1).

###### Proof.

Suppose that $Y$, $W$, and $H$ is an exact solution for the BANMF problem, and
that $\hat{W}$ and $\hat{H}$ are the Booleanized matrix of $W$ and $H$ with
threshold $\delta=0$. We show that
$X=\hat{W}\operatorname{\otimes}_{B}\hat{H}$.
First, following the logic in the proof of Theorem 1, we have that:

$(\hat{W}\operatorname{\otimes}_{B}\hat{H})_{ij}=\begin{cases}\begin{aligned}
&0&\text{ if }(WH)_{ij}=0\\\ &1&\text{ if
}(WH)_{ij}>0\end{aligned}\end{cases}$ (4)

Second, by the constraint of the optimization and the assumption that $Y=WH$,
we also have:

$X_{ij}=\begin{cases}\begin{aligned} &0&\text{ if }Y_{ij}=(WH)_{ij}=0\\\
&1&\text{ if }Y_{ij}=(WH)_{ij}>0\end{aligned}\end{cases}$ (5)

These two above facts imply that $X=\hat{W}\operatorname{\otimes}_{B}\hat{H}$
or that $\hat{W}$ and $\hat{H}$ is an exact solution for the BMF problem. ∎

###  3.2 Algorithm

Our implementation of the BANMF algorithm cycles between updating the NMF
factors $W$, $H$, and the corresponding auxiliary matrix $Y$. We use the
multiplicative update method introduced by Lee and Seung [6] for the $W$ and
$H$ updates, but any other alternating NMF algorithms are also suitable. For
the update of the auxiliary variable $Y$, we project $Y$ onto the feasible set
defined by the constraints. Algorithm 1 is an outline of a procedure to solve
the BANMF problem (2). To find the appropriate threshold parameters
$\delta_{W}$ and $\delta_{H}$, and the corresponding Boolean matrices,
$\hat{W}$ and $\hat{H}$, we employ Algorithm 2.

Input: Boolean matrix $X$, latent dimension $k$, maximum iteration $Niter$

Result: $Y,W,H$

Initialization: $W=rand(N,k)$, $H=rand(k,M)$, $Y=X$

while  _iter $\leq$ Niter_ do

(1) Update $W$ using NMF multiplicative update

$W_{ij}=W_{ij}\dfrac{(YH^{T})_{ij}}{(WHH^{T})_{ij}}$

(2) Update $H$ using NMF multiplicative update

$H_{ij}=H_{ij}\dfrac{(W^{T}Y)_{ij}}{(W^{T}WH)_{ij}}$

(3) Update $Y_{ij}$ for $(i,j)\in support(X)$

$Y_{ij}=\begin{cases}1\quad\text{if}\quad(WH)_{ij}<1\\\
(WH)_{ij}\quad\text{if}\quad 1\leq(WH)_{ij}\leq k\\\
k\quad\text{if}\quad(WH)_{ij}\geq k\end{cases}$

endwhile;

Algorithm 1 BANMF

Input: Boolean matrix $X$, nonnegative matrices $W$ and $H$, number of points
$npoint$

Result: $\hat{W},\hat{H}$

$W_{point}=\text{linspace(min(W), max(W), npoint)}$

$H_{point}=\text{linspace(min(H), max(H), npoint)}$

$\delta^{*}_{W},\delta^{*}_{H}=\underset{\delta_{W}\in W_{point},\delta_{H}\in
H_{point}}{\operatorname{argmin}}|X-\hat{W}^{\delta_{W}}\operatorname{\otimes}\hat{H}^{\delta_{H}}|$

$\hat{W}=\hat{W}^{\delta^{*}_{W}}$

$\hat{H}=\hat{H}^{\delta^{*}_{H}}$

Algorithm 2 Booleanization

####  3.2.1 Convergence

We now discuss the convergence of our implementation of the BANMF algorithm.
We show that Algorithm 1 results in a non-increasing updates to the objective
function:

######  Theorem 3 (Nonincreasing property of BANMF algorithm).

The objective function $||Y-WH||_{F}$ in the boolean auxiliary nonnegative
matrix factorization (2) is non-increasing under the BANMF update rules.

###### Proof.

Here we just need to show that
$||Y_{t-1}-W_{t-1}H_{t-1}||\geq||Y_{t}-W_{t}H_{t}||$ at an arbitrary iteration
$t$.
As shown in [6], the distance $||X-WH||_{F}$ is non-increasing under the
multiplicative update rules. Applying this to the updated $W_{t},H_{t}$ at the
iteration $t$, we have:

$||Y_{t-1}-W_{t-1}H_{t-1}||\geq||Y_{t-1}-W_{t}H_{t}||.$ (6)

The next step is updating the elements of $Y_{t-1}$ using the update rule in
step 3 of Algorithm 1. Note that the update rule minimizes the entry-wise
error of $|(Y_{t})_{ij}-(W_{t}H_{t})_{ij}|$ within the solution space for $Y$.
Indeed, each matrix entry for $Y_{t}$ either has zero error (for
$1\leq(WH)_{ij}\leq k)$ or has minimal error (for $(WH)_{ij}\leq 1$ and
$(WH)_{ij}\geq k$) which is an improvement over
$|(Y_{t-1})_{ij}-(W_{t}H_{t})_{ij}|$. Hence, this step produces $Y_{t}$ which
further reduces the error

$||Y_{t-1}-W_{t}H_{t}||\geq||Y_{t}-W_{t}H_{t}||$ (7)

From inequality (6) and (7), we get
$||Y_{t-1}-W_{t-1}H_{t-1}||\geq||Y_{t}-W_{t}H_{t}||$ as needed. Therefore, the
objective function of the problem (2) is non-increasing under BANMF update
rules. ∎

####  3.2.2 Regularized BANMF

Additionally, to aid in the Booleanization procedure, the factors $W$ and $H$
can be guided to be binary factors with an additional regularization during
the optimization. Here we incorporate the $x^{2}-x$ based regularization used
in [9, 15]. Our regularized problem is formed as follows:

$\displaystyle\underset{Y,W,H}{\text{minimize}}$
$\displaystyle||Y-WH||_{F}+\dfrac{1}{2}\lambda||W^{2}-W||^{2}_{F}$ (8)
$\displaystyle+\dfrac{1}{2}\lambda||H^{2}-H||^{2}_{F}$ subject to
$\displaystyle 1\leq Y_{i,j}\leq k,\text{ if }X_{i,j}=1$ $\displaystyle
Y_{i,j}=0,\text{ if }X_{i,j}=0$ $\displaystyle W\in\mathbb{R}_{+}^{N\times
k},$ $\displaystyle H\in\mathbb{R}_{+}^{k\times M}.$

Notice that under the exact solution, the regularization terms will be zeros.
Therefore, Theorem 1 and Theorem 2 are also true for the regularized problem
(8). With regularization, the update rules do not have the non-increasing
property anymore, but empirically we observe reasonable convergence in
practice. Figure 1 shows the convergence of the regularized BANMF on a
synthetic dataset. Our implementation of the regularized BANMF is described in
Algorithm 3.

Input: Boolean matrix $X$, latent dimension $k$, maximum iteration $Niter$,
regularization parameter $\lambda$

Result: $\hat{W}^{(\delta_{W})}$, $\hat{H}^{(\delta_{H})}$

Initialization: $W=rand(N,k)$, $H=rand(k,M)$, $Y=X$

while  _iter $\leq$ Niter_ do

(1) Update $W$ and $H$ using NMF multiplicative update

$W_{ij}=W_{ij}\dfrac{(YH^{T})_{ij}+3\lambda
W_{ij}^{2}}{(WHH^{T})_{ij}+2\lambda W_{ij}^{3}+\lambda W_{ij}^{2}}$

$H_{ij}=H_{ij}\dfrac{(W^{T}Y)_{ij}+3\lambda
H_{ij}^{2}}{(W^{T}WH)_{ij}+2\lambda H_{ij}^{3}+\lambda H_{ij}^{2}}$

(2) Update $Y_{ij}$ for $(i,j)\in support(X)$

$Y_{ij}=\begin{cases}1\quad\text{if}\quad(WH)_{ij}<1\\\
(WH)_{ij}\quad\text{if}\quad 1\leq(WH)_{ij}\leq k\\\
k\quad\text{if}\quad(WH)_{ij}\geq k\end{cases}$

endwhile;

Boolean algorithm is applied to identify $\delta_{W}$ and $\delta_{H}$ for
$\hat{W}^{(\delta_{W})}$, $\hat{H}^{(\delta_{H})}$

Algorithm 3 Regularized BANMF Figure 1: Convergence of regularized BANMF on a
synthetic dataset.

##  4 Experimental Results

In this section, we experimentally evaluate the performance of the BANMF
algorithm. For comparison we employ ASSO [8], NMF [6], and post nonlinear
penalty function (PNL-PF) [9] where the solutions are thresholded using
Booleanization Algorithm 2 to be binary. First, we demonstrate that if the
non-negative rank and the Boolean rank are different, then BANMF, and in
particular Regularized BANMF, outperform the other methods. Next, we
investigate how methods perform with regard to different noise and density
levels. We then show that BANMF outperforms its strongest contender, PNL-PF,
in execution time for random matrices. Lastly, the methods are evaluated when
applied to three real datasets: (1) UCI Zoo, (2) Congressional voting records,
and (3) Lung cancer with probes.

###  4.1 Generating Boolean synthetic data

First, we briefly describe our process for generating Boolean synthetic data.
For each experiment, we will generate random matrices

$X^{N\times M}=W^{M\times k}\operatorname{\otimes}_{B}H^{k\times m}$

where

$W_{ij},H_{ij}\sim Bernoulli(p).$

This stochastic process generates data with different features which
probabilistically depend on the parameters. In each experiment, this process
will be augmented to analyze statistical performance across various data
constraints.

To construct a Boolean rank $k$ data matrix with a desired density $d$, we
compute the probability that any given entry of $X$ is one,
$P(X_{ij}=1)=P(\lor_{l=1}^{k}W_{il}\land
H_{lj}=1)=1-(1-P(W_{il}=1)*P(H_{lj}=1))^{k}$. By imposing that the densities
of $W$ and $H$ are the same, and inverting the probability formula, we arrive
with the equation for the density of $W$ and $H$ that correspond to the
desired density of $X$, $\sqrt{1-(1-d)^{\frac{1}{k}}}$.

###  4.2 Different density levels

In the first experiment, we investigated the performance of the methods on
random Boolean matrices with different density levels. Here 100 different
$50\times 50$ matrices with $k=5$ were generated for each of three different
density levels: $20\%,50\%$ and $80\%$.

Figure 2 shows ASSO performs rather poorly across all density levels. For
sparse matrices, BANMF and REG-BANMF perform quite well. However for dense
matrices, PNLPF outperforms every method. Overall, (regularized) BANMF
consistently performs well across different sparsity levels.

Figure 2: Performance on synthetic datasets at different density levels.

###  4.3 Different noise levels

Figure 3: Mean and standard deviation of Boolean relative error of four
methods under three noise levels across 100 datasets. (Regularized) BANMF are
consistently outperform ASSO and PNL-PF.

The next experiment is to investigate the performance of all methods under the
effect of noise. Here 100 different $50\times 50$ matrices with $50\%$ density
and latent space $k=5$ are generated with no noise. Then Bernoulli noise is
added to the synthetic data according to different noise thresholds.
Concretely,

$X^{50\times 50}=W^{50\times 5}\operatorname{\otimes}_{B}H^{5\times 50}+_{f}E$

where

$W_{ij},H_{ij}\sim Bernoulli(p)\quad E_{i,j}\sim Bernoulli(p_{E})$

and the flipping operation $+_{f}$ is defined as

$X_{ij}=\begin{cases}X_{ij}\quad&\text{if }E_{ij}=0\\\ \neg
X_{ij}\quad&\text{if }E_{ij}=1.\\\ \end{cases}$ (9)

Each of the 100 matrices dataset is corrupted with noise levels ($0\%,1\%$ and
$5\%$). Figure 3 shows the mean and standard deviation of relative error
across the 100 datasets, in which regularized BANMF consistently outperform
ASSO and PNLPF. Moreover, BANMF outperforms PNLPF at lower noise thresholds.

###  4.4 BANMF and NMF on different-rank data

The NMF algorithm searches for a best nonnegative rank approximation to
$X=WH$. With BANMF being algorithmically similar to NMF, we are naturally
interested in the performance when there is a gap between the Boolean rank and
the nonnegative rank. Theoretically, it is known that the nonnegative rank is
greater than or equal to the Boolean rank [2]. To empirically investigate the
effect of a gap between the ranks, we generate a suite of Boolean matrices
where a gap between $\operatorname{rank}_{+}(X)$ and
$\operatorname{rank}_{\mathbb{B}}(X)$ exists. This is done as follows.

For each $N,M=10,11,\dots,50$, $k=2,3,\dots 6$, and density levels
$25\%,50\%,75\%$, we generate 5 Boolean matrices. We then want to measure the
difference $\operatorname{rank}_{+}(X)-\operatorname{rank}_{\mathbb{B}}(X)$
for each Boolean matrix $X$. Unfortunately, computing either the nonnegative
or Boolean rank is NP hard. However we can efficiently estimate a lower bound
on the rank gap. Indeed by construction,
$\operatorname{rank}_{\mathbb{B}}(X)\leq k$ and
$\operatorname{rank}(X)\leq\operatorname{rank}_{+}(X)$. Thus if
$k\leq\operatorname{rank}(X)$, then

$\operatorname{rank}_{+}(X)-\operatorname{rank}_{\mathbb{B}}(X)\geq\operatorname{rank}(X)-k.$

Thus for each generated matrix, we check that $k\leq\operatorname{rank}(X)$.
If this check fails, we generate a new Boolean matrix with the same
parameters. We then compute $\operatorname{rank}(X)-k$ to estimate the lower
bound on the gap.

Figure 4 depicts the average error of each method as a function of the lower
bound estimation of the rank gap. For lower rank gap, many of the methods are
comparable. However the error in using NMF grows as the gap between the
nonnegative and Boolean ranks grows, while the regularized BANMF algorithm is
constant and low, and the unregularized BANMF falls somewhere in between. The
regularized NMF’s error also grows with the rank gap, while PNLPF seems to
stay constant with a higher error level compared to regularized BANMF.

Figure 4: Performance of algorithms on Boolean/nonnegative rank gap dataset.

###  4.5 Execution time

Since PNLPF’s performance is comparable with BANMF in some previous cases, we
further compare the complexity between PNLPF and BANMF. Given the high level
of nonlinearity of PNLPF, it is expected to be more computationally expensive
than BANMF algorithm. In the first comparison, the execution times of 1000
iterations for datasets of different data sizes are measured. Figure 5 shows
the execution time in log-scale of PNLPF and BANMF. The results show that when
the data size is getting larger, PNLPF is increasingly more expensive. For
example, at the data size of $500\times 500$, the execution time of PNLPF is
about $2^{4}=16$ times longer than of BANMF. In the second comparison, Figure
6 shows the time that each method converges in Boolean error rate for
different datasets. The result shows that as the data size increases, PNLPF
takes a longer time to converge. In the conclusion, BANMF is less
computationally expensive compared to PNLPF.

Figure 5: Execution time in log-scale of BANMF and PNLPF for 5 different size
datasets. Overall, PNLPF is more expensive than BANMF, and even more expensive
when the data size is larger. Figure 6: Convergence time in Boolean error rate
of BANMF and PNLPF. As the data is larger, PNLPF takes increasingly longer
time to converge.

###  4.6 Application in real data

In this section, we demonstrate the performance of (regularized) BANMF and the
other methods on three real datasets: UCI Zoo, Congressional voting records
and Lung cancer with probes LUCAP0.

####  4.6.1 UCI zoo dataset

We first study the performance of BNMF-MC on decomposing the UCI zoo dataset
111https://archive.ics.uci.edu/ml/datasets/Zoo [3]. This dataset consists of
101 animals and 14 binary features, which results in a $101\times 14$ Boolean
matrix. Figure 7 shows that regularized BANMF and NMF have the best
performance, especially with higher latent dimensions. Furthermore, Figure 8
shows the original data and reconstructed data from rank-3 decompositions. All
methods seem to capture the Boolean structure of the data.

Figure 7: UCI Zoo dataset. Regularized BANMF and NMF have the best
performances, especially with higher latent dimensions. Figure 8: Original
data versus reconstructed data from a rank-3 decomposition of (regularized)
BANMF and other methods. All methods seem to capture the Boolean structure of
the data.

####  4.6.2 Congressional voting records dataset

The second experiment uses the Congressional Voting Records Dataset from UCI
dataset [3]. This data includes the votes of the U.S. House of Representatives
Congressmen on the 16 key votes. After removing samples with missing data, the
Boolean data matrix has the dimension of $232\times 16$. Figure 9 shows that
(regularized) NMF are better with (regularized) BANMF are slightly behind.

Figure 9: Congressional voting records dataset - (regularized) NMF has the
lowest errors with (regularized) BANMF are slightly behind.

####  4.6.3 Lung cancer with probes dataset (LUCAP0)

The third data set is the LUCAP0 (Lung Cancer with Probes) data which contains
toy data generated artificially by causal Bayesian networks with binary
variables, which being used in [9]. This model is a medical application for
diagnosing lung cancer.222http://www.causality.inf.ethz.ch/data/LUCAS.html.
The dimensions of the data are $2000\times 143$. For this dataset, Figure 10
shows that (regularized) NMF are slightly better here. Additionally, PNLPF
performs well with low latent dimension, but is outperformed by regularized
BANMF when the latent dimension gets larger. This PNLPF’s behavior can also be
observed from previous datasets. Moreover, given this size of the data, PNLPF
is more computationally expensive than BANMF and NMF (see subsection 4.5 for
the execution time comparison.)

Figure 10: Lung cancer with probes dataset. (Regularized) NMF has the lowest
error with higher latent dimension.

##  5 Discussion

In this paper, we propose an approach to solve the Boolean matrix
factorization problem using a nonnegative optimization with the constraint
over an auxiliary matrix. With this constraint, the algorithm can overcome the
difference between Boolean algebra and normal algebra to capture the Boolean
structure of the data. The solution is then thresholded to give a solution for
BMF problem. We show that the exact solution spaces of the two problem are
equivalent, and that the BANMF algorithm has the nonincreasing property. The
(regularized) BANMF algorithm is shown to be on outperform other methods in
experiments on synthetic examples at particular density, and noisy regimes.
Furthermore, these algorithms prove to be superior methods when there is a gap
between Boolean rank and nonnegative rank. Outside these regimes, the
(regularized) BANMF are still competitive with other state of the art
algorithms.

## References

  * [1] Radim Belohlavek and Vilem Vychodil.  Discovery of optimal factors in binary data via a novel method of matrix decomposition.  Journal of Computer and System Sciences, 76(1):3–20, 2010.
  * [2] Derek DeSantis, Erik Skau, and Boian Alexandrov.  Factorizations of binary matrices–rank relations and the uniqueness of boolean decompositions.  arXiv preprint arXiv:2012.10496, 2020.
  * [3] Dheeru Dua and Casey Graff.  UCI machine learning repository, 2017.
  * [4] B Everett.  An introduction to latent variable models.  Springer Science & Business Media, 2013.
  * [5] Daniel D Lee and H Sebastian Seung.  Learning the parts of objects by non-negative matrix factorization.  Nature, 401(6755):788–791, 1999.
  * [6] Daniel D Lee and H Sebastian Seung.  Algorithms for non-negative matrix factorization.  In Advances in neural information processing systems, pages 556–562, 2001.
  * [7] Claudio Lucchese, Salvatore Orlando, and Raffaele Perego.  Mining top-k patterns from binary datasets in presence of noise.  In Proceedings of the 2010 SIAM International Conference on Data Mining, pages 165–176. SIAM, 2010.
  * [8] Pauli Miettinen, Taneli Mielikäinen, Aristides Gionis, Gautam Das, and Heikki Mannila.  The discrete basis problem.  IEEE transactions on knowledge and data engineering, 20(10):1348–1362, 2008.
  * [9] Sebastian Miron, Mamadou Diop, Anthony Larue, Eddy Robin, and David Brie.  Boolean decomposition of binary matrices using a post-nonlinear mixture approach.  Signal Processing, 178:107809, 2021.
  * [10] S. D. Monson, N. J. Pullman, and R. Rees.  A Survey of Clique and Biclique Coverings and Factorizations of (0,1)-Matrices.  Bull. ICA, 14:17–86, 1995.
  * [11] Siamak Ravanbakhsh, Barnabás Póczos, and Russell Greiner.  Boolean matrix factorization and noisy completion via message passing.  In ICML, volume 69, pages 945–954, 2016.
  * [12] Tammo Rukat, Chris C Holmes, Michalis K Titsias, and Christopher Yau.  Bayesian boolean matrix factorisation.  arXiv preprint arXiv:1702.06166, 2017.
  * [13] Charles Spearman.  “General intelligence,” objectively determined and measured.  The American Journal of Psychology, 15, 1961.
  * [14] G. W. Stewart.  On the early history of the singular value decomposition.  SIAM review, 35(4):551–566, 1993.
  * [15] Zhongyuan Zhang, Tao Li, Chris Ding, and Xiangsun Zhang.  Binary matrix factorization with applications.  In Seventh IEEE international conference on data mining (ICDM 2007), pages 391–400. IEEE, 2007.


RobustNeuralNetworks.jl: a Package for Machine Learning and Data-Driven Control with Certified Robustness

1]Nicholas H. Barbara
1]Max Revay
1]Ruigang Wang
1]Jing Cheng
1]Ian R. Manchester
[1]University of Sydney, Australian Centre for Robotics

Neural networks are typically sensitive to small input perturbations, leading to unexpected or brittle behaviour. We present : a Julia package for neural network models that are constructed to naturally satisfy a set of user-defined robustness constraints. The package is based on the recently proposed Recurrent Equilibrium Network (REN) and Lipschitz-Bounded Deep Network (LBDN) model classes, and is designed to interface directly with Julia's most widely-used machine learning package, . We discuss the theory behind our model parameterization, give an overview of the package, and provide a tutorial demonstrating its use in image classification, reinforcement learning, and nonlinear state-observer design.

§ INTRODUCTION

Modern machine learning relies heavily on rapidly training and evaluating neural networks in problems ranging from image classification <cit.> to robotic control <cit.>. Most neural network architectures have no robustness certificates, and can be sensitive to adversarial attacks and other input perturbations <cit.>. Many that address this brittle behaviour rely on explicitly enforcing constraints during training to smooth or stabilize the network response <cit.>. While effective on small-scale problems, these methods are computationally expensive, making them slow and difficult to scale up to complex real-world problems.

Recently, we proposed the Recurrent Equilibrium Network (REN) <cit.> and Lipschitz-Bounded Deep Network (LBDN) or sandwich layer <cit.> model classes as computationally efficient solutions to these problems. The REN architecture is flexible in that it includes all common neural network models, such as multi-layer-perceptrons (MLPs), convolutional neural networks (CNNs), and recurrent neural networks (RNNs). The weights and biases in RENs are directly parameterized to naturally satisfy behavioural constraints chosen by the user. For example, we can build a REN with a given Lipschitz constant to ensure its output is quantifiably less sensitive to input perturbations. LBDNs are specializations of RENs with the specific feed-forward structure of deep neural networks like MLPs or CNNs and built-in guarantees on the Lipschitz bound.

The direct parameterization of RENs and LBDNs means that we can train models with standard, unconstrained optimization methods (such as stochastic gradient descent) while also guaranteeing their robustness. Achieving the “best of both worlds” in this way is the main advantage of the REN and LBDN model classes, and allows the user to freely train robust models for many common machine learning problems, as well as for more challenging real-world applications where safety is critical.

This papers presents : a package for neural network models that naturally satisfy robustness constraints. The package contains implementations of the REN and LBDN model classes introduced in <cit.> and <cit.>, respectively, and relies heavily on key features of the Julia language <cit.> (such as multiple dispatch) for an efficient implementation of these models. The purpose of  is to make our recent research in robust machine learning easily accessible to users in the scientific and machine learning communities. With this in mind, we have designed the package to interface directly with  <cit.>, Julia's most widely-used machine learning package, making it straightforward to incorporate robust neural networks into existing Julia code.

The paper is structured as follows. Section <ref> provides an overview of the  package, including a brief introduction to the model classes (Sec. <ref>), their robustness certificates (Sec. <ref>), and their implementation (Sec. <ref>). Section <ref> guides the reader through a tutorial with three examples to demonstrate the use of RENs and LBDNs in machine learning: image classification (Sec. <ref>), reinforcement learning (Sec. <ref>), and nonlinear state-observer design (Sec. <ref>). Section <ref> offers some concluding remarks and future directions for robust machine learning with . For more detail on the theory behind RENs and LBDNs, and for examples comparing their performance to current state-of-the-art methods on a range of problems, we refer the reader to <cit.> and <cit.> (respectively).

§ PACKAGE OVERVIEW

contains two classes of neural network models: RENs and LBDNs. This section gives a brief overview of the two model architectures and how they are parameterized to automatically satisfy robustness certificates. We also provide some background on the different types of robustness metrics used to construct the models.

§.§ What are RENs and LBDNs?

A Lipschitz-Bounded Deep Network (LBDN) is a (memoryless) deep neural network model with a built-in upper-bound on its Lipschitz constant (Sec. <ref>). Suppose the network has inputs $u \in\mathbb{R}^{n_u}$, outputs $y \in \mathbb{R}^{n_y}$, and hidden units $z_k \in \mathbb{R}^{n_k}$. We define LBDNs as $L$-layer feed-forward networks, much like MLPs or CNNs, described by the following recursive equations.

\begin{align}
z_0 &= x \label{eqn:lbdn-z0}\\
z_{k+1} &= \sigma(W_k z_k + b_k), \quad k = 0, \ldots, L-1 \label{eqn:lbdn-sandwich}\\
y &= W_L z_L + b_L \label{eqn:lbdn-output}
\end{align}

A Recurrent Equilibrium Network (REN) is a neural network model described by linear system in feedback with a nonlinear activation function. Denote $x_t \in \mathbb{R}^{n_x}$ as the internal states of the system, $u_t \in\mathbb{R}^{n_u}$ as its inputs, and $y_t \in \mathbb{R}^{n_y}$ as its outputs. Mathematically, a REN can be written as

\begin{align}
\begin{bmatrix}
x_{t+1} \\ v_t \\ y_t
\end{bmatrix}&=
\overset{W}{\overbrace{
\left[
\begin{array}{c|cc}
A & B_1 & B_2 \\ \hline
C_{1} & D_{11} & D_{12} \\
C_{2} & D_{21} & D_{22}
\end{array}
\right]
\begin{bmatrix}
x_t \\ w_t \\ u_t
\end{bmatrix}+
\overset{b}{\overbrace{
\begin{bmatrix}
b_x \\ b_v \\ b_y
\end{bmatrix}
}}, \label{eqn:ren-G}\\
\sigma(v_{t}^1) & \sigma(v_{t}^2) & \cdots & \sigma(v_{t}^q)
\end{bmatrix}^\top, \label{eqn:ren-sigma}
\end{align}

where $v_t, w_t \in \mathbb{R}^{n_v}$ are the inputs and outputs of neurons and $\sigma$ is the activation function. Graphically, this is equivalent to Figure <ref>, where the linear (actually affine) system $G$ is given by Equation <ref>.

Feedback structure of a recurrent equilibrium network.

See <cit.> and <cit.> for more on RENs and LBDNs, respectively.

<cit.> makes special mention of “acyclic” RENs, which have a lower-triangular $D_{11}$ matrix. These are significantly more efficient to evaluate and train than a REN with dense $D_{11}$ matrix, and the performance is similar across a range of problems. All RENs in the current implementation of  are therefore acyclic RENs.

§.§ Robust metrics and IQCs

All neural network models in  are designed to satisfy a set of user-defined robustness constraints. There are a number of different robustness criteria which our RENs can satisfy. Some relate to the internal dynamics of the model, others relate to the input-output map. LBDNs are less general, and are specifically constructed to satisfy Lipschitz bounds (Sec. <ref>).

§.§.§ Contracting systems

Firstly, all of our RENs are contracting systems. This means that they exponentially “forget” their initial conditions. If the system starts at two different initial conditions but is given the same input sequence, the internal states will exponentially converge over time. Figure <ref> shows an example of a contracting REN with one input and a single internal state, where two simulations of the system start with different initial conditions but are provided the same sinusoidal input. See <cit.> for a detailed introduction to contraction theory for dynamical systems.

Simulation of a contracting REN with a single internal state. The system is simulated from two different initial states with the same sinusoidal input. The contracting system exponentially forgets its initial condition.

§.§.§ Incremental IQCs

We define additional robustness criteria on the input-output map of RENs with incremental integral quadratic constraints (IQCs). Suppose we have a model $\mathcal{M}$ starting at two different initial conditions $a,b$ with two different input signals $u, v$, and consider their corresponding output trajectories $y^a = \mathcal{M}_a(u)$ and $y^b = \mathcal{M}_b(v).$ The model $\mathcal{M}$ satisfies the IQC defined by matrices $(Q, S, R)$ if

\begin{equation}
\sum_{t=0}^T
\begin{bmatrix}
y^a_t - y^b_t \\ u_t - v_t
\end{bmatrix}^\top
\begin{bmatrix}
Q & S^\top \\ S & R
\end{bmatrix}
\begin{bmatrix}
y^a_t - y^b_t \\ u_t - v_t
\end{bmatrix}
\ge -d(a,b)
\quad \forall \, T
\end{equation}

for some function $d(a,b) \ge 0$ with $d(a,a) = 0$, where $0 \preceq Q \in \mathbb{R}^{n_y\times n_y}$, $S\in\mathbb{R}^{n_u\times n_y},$ $R=R^\top \in \mathbb{R}^{n_u\times n_u}.$

In general, the IQC matrices $(Q,S,R)$ can be chosen (or optimized) to meet a range of performance criteria. There are a few special cases that are worth noting.

§.§.§ Lipschitz bounds (smoothness)

If $Q = -\frac{1}{\gamma}I$, $R = \gamma I$, $S = 0$ for some $\gamma \in \mathbb{R}$ with $\gamma \ne 0$, the model $\mathcal{M}$ satisfies a Lipschitz bound (incremental $\ell^2$-gain bound) of $\gamma$, defined by

\begin{equation}
\|\mathcal{M}_a(u) - \mathcal{M}_b(v)\|_T^2 \le \gamma^2 \|u - v\|_T^2.
\end{equation}

Qualitatively, the Lipschitz bound is a measure of how smooth the network is. If the Lipschitz bound $\gamma$ is small, then small changes in the inputs $u,v$ will lead to small changes in the model output. If $\gamma$ is large, then the model output might change significantly for even small changes to the inputs. This can make the model more sensitive to noise, adversarial attacks, and other input disturbances.

As the name suggests, all LBDN models are all constructed to have a user-tunable (or learnable) Lipschitz bound.

§.§.§ Incremental passivity

We have implemented two versions of incremental passivity. In each case, the network must have the same number of inputs and outputs.

* If $Q = 0, R = -2\nu I, S = I$ where $\nu \ge 0$, the model is incrementally passive (incrementally strictly input passive if $\nu > 0$). Mathematically, the following inequality holds.

\begin{equation}
\langle \mathcal{M}_a(u) - \mathcal{M}_b(v), u-v \rangle_T \ge \nu \| u-v\|^2_T
\end{equation}

* If $Q = -2\rho I, R = 0, S = I$ where $\rho > 0$, the model is incrementally strictly output passive. Mathematically, the following inequality holds.

\begin{equation}
\langle \mathcal{M}_a(u) - \mathcal{M}_b(v), u-v \rangle_T \ge \rho \| \mathcal{M}_a(u) - \mathcal{M}_b(v)\|^2_T
\end{equation}

For more details on IQCs and their use in RENs, please see <cit.>.

§.§ Direct and explicit parameterizations

The key advantage of the models in  is that they naturally satisfy the robustness constraints of Section <ref> – i.e., robustness is guaranteed by construction. There is no need to impose additional (possibly computationally-expensive) constraints while training a REN or an LBDN. One can simply use unconstrained optimization methods like gradient descent and be sure that the final model will satisfy the robustness requirements.

We achieve this by constructing the weight matrices and bias vectors in our models to automatically satisfy specific linear matrix inequalities (see <cit.> for details). The learnable parameters of a REN or LBDN are a set of free, unconstrained variables $\theta \in \mathbb{R}^N$. When the set of learnable parameters is exactly $\mathbb{R}^N$, we call this a direct parameterization. Equations <ref> to <ref> describe the explicit parameterizations of RENs and LBDNs: model structures that can be called and evaluated on data. For a REN, the explicit parameters are $\bar{\theta} := [W, b]$, and for an LBDN they are $\bar{\theta} := [W_0, b_0, \ldots, W_L, b_L]$. The mapping $\theta \mapsto \bar{\theta}$ depends on the specific robustness constraints to be imposed on the explicit model.

§.§.§ Implementation

RENs are defined by two abstract types in . Subtypes of  hold all the information required to directly parameterize a REN satisfying some robustness properties. For example, to initialise the direct parameters of a contracting REN with 1 input, 10 states, 20 neurons, 1 output, and a  activation function, we use the following. The direct parameters $\theta$ are stored in .

[language = Julia]
using Flux, RobustNeuralNetworks

T  = Float32
nu, nx, nv, ny = 1, 10, 20, 1
model_ps = ContractingRENParamsT(
nu, nx, nv, ny; nl=Flux.relu)

println(model_ps.direct) # Access direct params

Subtypes of  represent RENs in their explicit form which can be evaluated on data. The conversion from direct to explicit parameters $\theta \mapsto \bar{\theta}$ is performed when the REN is constructed and the explicit parameters $\bar{\theta}$ are stored in .

[language = Julia]
model = REN(ps)          # Create explicit model
println(model.explicit)  # Access explicit params

Figure <ref> illustrates this architecture. We use a similar interface based on  and  for LBDNs.

Association of models and their parameters in (a)  and (b) . In (a), model parameters $\theta$ are associated with the . In (b), the direct parameters $\theta$ are associated with the parameterization , and are converted to explicit parameters $\bar{\theta}$ when the  is constructed with .

§.§.§ Types of direct parameterizations

There are currently four REN parameterizations implemented in this package:

*  parameterizes contracting RENs with a user-defined upper bound on the contraction rate.

*  parameterizes RENs with a user-defined (or learnable) Lipschitz bound of $\gamma \in (0,\infty)$.

*  parameterizes incrementally input passive RENs with user-tunable passivity parameter $\nu \ge 0$.

*  parameterizes RENs satisfying some general behavioural constraints defined by an incremental IQC with parameters (Q,S,R).

There is currently only one LBDN parameterization implemented in :

*  parameterizes dense (fully-connected) LBDNs with a user-defined or learnable Lipschitz bound. A dense LBDN is effectively a Lipschitz-bounded MLP.

We intend to add  to parameterize the convolutional LBDNs in <cit.> in future iterations of the package.

§.§.§ Explicit model wrappers

When training a REN or LBDN, we learn and update the direct parameters $\theta$ and convert them to the explicit parameters $\bar{\theta}$ only for model evaluation. The main constructors for explicit models are  and .

Users familiar with  will be used to creating a model once and then training it on their data. The typical workflow is as follows.

[language = Julia]
using Flux

# Define a model and a loss function
model = Flux.Chain(
Flux.Dense(1 => 10, Flux.relu),
Flux.Dense(10 => 1, Flux.relu)

loss(model, x, y) = Flux.mse(model(x), y)

# Training data of 20 batches
T = Float32
xs, ys = rand(T,1,20), rand(T,1,20)
data = [(xs, ys)]

# Train the model for 50 epochs
opt_state = Flux.setup(Adam(0.01), model)
for _ in 1:50
Flux.train!(loss, model, data, opt_state)

When training a model constructed from  or , we need to back-propagate through the mapping from direct (learnable) parameters to the explicit model. We must therefore include the model construction as part of the loss function. If we do not, then the auto-differentiation engine has no knowledge of how the model parameters affect the loss, and will return zero gradients. Here is an example with an , where the  is defined by the direct parameterization stored in .

[language = Julia]
using Flux, RobustNeuralNetworks

# Define model parameterization and loss function
T = Float32
model_ps = DenseLBDNParamsT(1, [10], 1; nl=relu)
function loss(model_ps, x, y)
model = LBDN(model_ps)
Flux.mse(model(x), y)

# Training data of 20 batches
xs, ys = rand(T,1,20), rand(T,1,20)
data = [(xs, ys)]

# Train the model for 50 epochs
opt_state = Flux.setup(Adam(0.01), model_ps)
for _ in 1:50
Flux.train!(loss, model_ps, data, opt_state)

§.§.§ Separating parameters and models

For the sake of convenience, we have included the model wrappers  and  as alternatives to  and , respectively. These wrappers compute the explicit parameters each time the model is called rather than just once when they are constructed. Any model created with these wrappers can therefore be used exactly the same way as a regular  model, and there is no need for model construction in the loss function. One can simply replace the definition of the  model in the first code cell above with
[language = Julia]
model_ps = DenseLBDNParamsT(1, [10], 1; nl=relu)
model = DiffLBDN(model_ps)

and train the LBDN just like any other  model. We use these wrappers in Sections <ref> and <ref>.

We must, however, caution against their use in problems where a model is evaluated many times before being updated. The  and  are kept separate with  and  to offer flexibility. The main computational bottleneck in training a REN or LBDN is converting from the direct to explicit parameters (mapping $\bar{f}: \theta \mapsto \bar{\theta}$). For both model structures, this process involves a matrix inverse, where the number of matrix elements scales quadratically with the dimension of the model in a REN and the dimension of each layer in an LBDN (see <cit.>).

As per Section <ref>, direct parameters are stored in , while explicit parameters are computed when the  is created and are stored within it. In some applications (eg: reinforcement learning), a model is called many times with the same explicit parameters $\bar{\theta}$ before its learnable parameters $\theta$ are updated. It is therefore significantly more efficient to store the explicit parameters, use them many times, and then update them only when the learnable parameters change. We cannot store the direct and explicit parameters in the same  object since auto-differentiation in Julia does not support array mutation <cit.>. Instead, we separate the two.

The model wrappers  and  re-compute the explicit parameters every time the model is called. In applications where the learnable parameters are updated after one model call (eg: image classification), it is often more convenient and equally fast to use these wrappers.

In applications where the model is called many times before updating it (eg: reinforcement learning), use  or . They compute the explicit model when constructed and store it for later use, making them more efficient.

The computational benefits of separating models from their parameterizations is explored numerically in Section <ref>.

§ EXAMPLES

This section guides the reader through a set of examples to demonstrate how to use  for machine learning in Julia. We will consider three examples: image classification, reinforcement learning, and nonlinear observer design. These examples will provide further insight into the benefits of using robust models and the reasoning behind key design decisions made when developing the package.

We use  activation functions in all examples, but other choices of activation function (e.g: ) are equally valid. We note that any activation function used in a REN or LBDN must have a maximum slope of 1.0, as outlined in <cit.>. For more examples with RENs and LBDNs, please see the package documentation[<https://acfr.github.io/RobustNeuralNetworks.jl/>].

§.§ Image classification

Our first example features an LBDN trained to classify the MNIST dataset <cit.>. We will use this example to demonstrate how training image classifiers with LBDNs makes them robust to adversarial attacks thanks to the built-in Lipschitz bound. For a detailed investigation of the effect of Lipschitz bounds on classification robustness and reliability, please see <cit.>.

§.§.§ Load the data

We begin by loading the training and test data. [<https://juliaml.github.io/MLDatasets.jl/>] contains a number of common machine-learning datasets, including the MNIST dataset. To load the full dataset of 60,000 training images and 10,000 test images, one would run the following code.

[language = Julia]
using MLDatasets: MNIST

T = Float32
x_train, y_train = MNIST(T, split=:train)[:]
x_test,  y_test  = MNIST(T, split=:test)[:]

The feature matrices  and  are three-dimensional arrays where each 28x28 layer contains pixel data for a single handwritten number from 0 to 9 (eg: see Fig. <ref>). The labels  and  are vectors containing the classification of each image as a number from 0 to 9. We convert each of these to a format better suited to training with .

[language = Julia]
using Flux

# Reshape features for model input
x_train = Flux.flatten(x_train)
x_test  = Flux.flatten(x_test)

# Encode categorical outputs and store
y_train = Flux.onehotbatch(y_train, 0:9)
y_test  = Flux.onehotbatch(y_test,  0:9)
data = [(x_train, y_train)]

Features are now stored in a $28^2\times N$  where each column contains pixel data from a single image, and the labels have been converted to a $10\times N$  where each column contains a 1 in the row corresponding to the image's classification (eg: row 3 for an image showing the number 2) and a 0 otherwise.

§.§.§ Define a model

We can now construct an LBDN model to train on the MNIST dataset. The larger the model, the better the classification accuracy will be, at the cost of longer training times. The smaller the Lipschitz bound $\gamma$, the more robust the model will be to input perturbations (such as noise in the image). If $\gamma$ is too small, however, it can restrict the model flexibility and limit the achievable performance <cit.>. For this example, we use a small network of two 64-neuron hidden layers and set a Lipschitz bound of $\gamma=5.0$ just to demonstrate the method.

[language = Julia]
using RobustNeuralNetworks

# Model specification
nu = 28*28              # Inputs (size of image)
ny = 10                 # Outputs (classifications)
nh = fill(64,2)         # Hidden layers
γ  = 5.0f0              # Lipschitz bound 5.0

# Define parameters,create model
model_ps = DenseLBDNParamsT(nu, nh, ny, γ)
model = Chain(DiffLBDN(model_ps), Flux.softmax)

The  consists of two parts. The first is a callable  model constructed from its direct parameterization, which is defined by an instance of  as per Section <ref>. The output is then converted to a probability distribution using a  layer. Note that all  models can be combined with traditional neural network layers using .

We could also have constructed the  from a combination of  layers. Introduced in <cit.>, the  layer is a fully-connected or dense layer with a guaranteed Lipschitz bound of 1.0. We have designed the user interface for  to be as similar to that of  as possible.
[language = Julia]
model = Chain(
(x) -> (sqrt(γ) * x),
SandwichFC(nu => nh[1], relu; T),
SandwichFC(nh[1] => nh[2], relu; T),
(x) -> (sqrt(γ) * x),
SandwichFC(nh[2] => ny; output_layer=true, T),

This model is equivalent to a dense LBDN constructed with  or . We have included it as a convenience for users familiar with layer-wise network construction in , and recommend using it interchangeably with .

§.§.§ Define a loss function

A typical loss function for training on datasets with discrete labels is the cross entropy loss. We can use the  loss function shipped with .

[language = Julia]
loss(model,x,y) = Flux.crossentropy(model(x), y)

§.§.§ Train the model

We train the model over 600 epochs using two learning rates:  for the first 300, and  for the last 300. We use the  optimizer <cit.> and the default  method for convenience. Note that the  method updates the learnable parameters each time the model is evaluated on a batch of data, hence our choice of  as a model wrapper.

[language = Julia]
# Hyperparameters
epochs = 300
lrs = [1e-3,1e-4]

# Train with the Adam optimizer
for k in eachindex(lrs)
opt_state = Flux.setup(Adam(lrs[k]), model)
for i in 1:epochs
Flux.train!(loss, model, data, opt_state)

§.§.§ Evaluate the trained model

Our final model achieves training and test accuracies of approximately 99% and 98%, respectively, as shown in Table <ref>. We could improve this further by switching to a convolutional LBDN, as in <cit.>. Some examples of classifications given by the trained LBDN model are presented in Figure <ref>.

Examples of classifications from the trained LBDN model on the MNIST dataset.

§.§.§ Investigate robustness

The main advantage of using an LBDN for image classification is its built-in robustness to noise (or attacks) added to the image. This robustness is a direct benefit of the Lipschitz bound. As outlined in the Section <ref>, the Lipschitz bound effectively defines how “smooth” the network is: the smaller the Lipschitz bound, the less the network outputs will change as the inputs vary. For example, small amounts of noise added to the image will be less likely to change its classification. A detailed investigation into this effect is presented in <cit.>.

We can demonstrate the robustness of LBDNs by comparing the model to a standard MLP built from  layers. We first create a  network with the same layer structure as the LBDN.
[language = Julia]
# Initialisation functions
init = Flux.glorot_normal
initb(n) = Flux.glorot_normal(n)

# Build a dense model
dense = Chain(
Dense(nu, nh[1], relu;
init, bias=initb(nh[1])),
Dense(nh[1], nh[2], relu;
init, bias=initb(nh[2])),
Dense(nh[2], ny; init, bias=initb(ny)),

Training the  model with the same training loop used for the LBDN model results in a model that achieves training and test accuracies of approximately 98% and 97%, respectively, as shown in Table <ref>.

Training and test accuracy for the LBDN and Dense models on the MNIST dataset without perturbations.
Model structure     Training accuracy (%)     Test accuracy (%)
LBDN     98.7     97.6
Dense     98.4     96.9

As a simple test of robustness, we add uniformly-sampled random noise in the range $[-\epsilon, \epsilon]$ to the pixel data in the test dataset for a range of noise magnitudes $\epsilon \in [0, 200/255].$ We record the test accuracy for each perturbation size and store it for plotting.
[language = Julia]
using Statistics

# Get test accuracy as we add noise
uniform(x) = 2*rand(T, size(x)...) .- 1
compare(y, yh) =
maximum(yh, dims=1) .== maximum(y.*yh, dims=1)
accuracy(model, x, y) = mean(compare(y, model(x)))

function noisy_test_error(model, ϵ)
noisy_xtest = x_test .+ ϵ*uniform(x_test)
accuracy(model, noisy_xtest,  y_test)*100

ϵs = T.(LinRange(0, 200, 10)) ./ 255
lbdn_error  = noisy_test_error.((model,), ϵs)
dense_error = noisy_test_error.((dense,), ϵs)

Plotting the results in Figure <ref> very clearly shows that the  network, which has no guarantees on its Lipschitz bound, quickly loses its accuracy as small amounts of noise are added to the image. In contrast, the LBDN  maintains its accuracy even when the (maximum) perturbation size is as much as 80% of the maximum pixel values. This is an illustration of why image classification is one of the most promising use-cases for LBDN models. For a more detailed comparison of LBDN with state-of-the-art image classification methods, see <cit.>.

Comparison of test accuracy on the MNIST dataset as a function of random perturbation magnitude $\epsilon$. The LBDN model is significantly more robust than a standard  network.

§.§ Reinforcement learning

One of the original motivations for developing the model structures in   was to guarantee stability and robustness in learning-based control. Recently, we have shown that with a controller architecture based on a nonlinear version of classical Youla-Kucera parameterization <cit.>, one can learn over the space of all stabilising controllers for linear or nonlinear systems using standard reinforcement learning techniques, so long as our control policy is parameterized by a contracting, Lipschitz-bounded REN  <cit.>. This is an exciting result when considering learning controllers in safety-critical systems, such as in robotics.

In this example, we will demonstrate how to train an LBDN controller with reinforcement learning (RL) for a simple nonlinear dynamical system. This controller will not have any stability guarantees. The purpose of this example is simply to showcase the steps required to set up RL experiments for more complex systems with RENs and LBDNs.

§.§.§ Overview

Consider the simple mechanical system shown in Figure <ref>: a box of mass $m$ sits in a tub of fluid, held between the walls by two springs each with spring constant $k/2.$ The box can be pushed with a force $u.$ Its dynamics are

\begin{equation} \label{eqn:box-dynamics}
m\ddot{q} = u - kq - \mu \dot{q}^2
\end{equation}

where $\mu$ is the viscous damping coefficient due to the box moving through the fluid, and $\dot{q},\ddot{q}$ denote the velocity and acceleration of the box, respectively.

Mechanical system to be controlled. A box sits in a tub of fluid, suspended between two springs, and can be pushed by a force $u$ to different horizontal positions $q$.

We can write this as a (nonlinear) state-space model with state $x = (q,\dot{q})^\top,$ control input $u,$ and dynamics

\begin{equation}
\dot{x} = f(x,u) := \begin{bmatrix}
\dot{q} \\ (u - kq - \mu \dot{q}^2)/m
\end{bmatrix}.
\end{equation}

This is a continuous-time model of the dynamics. For our purposes, we need a discrete-time model. We can discretize the dynamics using a forward Euler approximation to get

\begin{equation}
x_{t+1} = f_d(x_t,u_t) := x_t + \Delta t \cdot f(x_t, u_t)
\end{equation}

where $\Delta t$ is the time-step. This approximation typically requires a small time-step for numerical stability, but is sufficient for our simple example. If physical accuracy was of concern, one could use a fourth (or higher) order Runge-Kutta scheme.

Our aim is to learn a controller $u = \mathcal{K}_\theta(x, q_\mathrm{ref}),$ defined by some learnable parameters $\theta,$ that can push the box to any goal position $q_\mathrm{ref}$ that we choose. Specifically, we want the box to:

* reach a (stationary) goal position $q_\mathrm{ref}$

* within a time period $T$.

The force required to keep the box at a static equilibrium position $q_\mathrm{ref}$ is $u_\mathrm{ref} = k q_\mathrm{ref}$ from Equation <ref>. We can encode these objectives into a cost function $J_\theta$ and write our RL problem as

\begin{equation} \label{eqn:rl-costfunc}
\min_\theta \mathbb{E} \left[ J_\theta \right],
\quad
J_\theta = \sum_{t=0}^{T-1} c_1 (\Delta q_t)^2 + c_2 \dot{q}_t^2 + c_3 (\Delta u_t)^2
\end{equation}

where $\Delta q_t = q_t - q_\mathrm{ref}$, $\Delta u_t = u_t - u_\mathrm{ref}$, $c_1, c_2, c_3$ are cost function weights, and the expectation is over different initial and goal positions of the box.

§.§.§ Problem setup

We start by defining the properties of our system and translating the dynamics into Julia code. For this example, we consider a box of mass $m=1$, spring constants $k=5,$ and a viscous damping coefficient $\mu = 0.5$. We will simulate the system over $T = 4$ s time horizons with a time-step of $\Delta t = 0.02$ s.

[language = Julia]
m = 1                   # Mass (kg)
k = 5                   # Spring constant (N/m)
μ = 0.5                 # Viscous damping (kg/m)

Tmax = 4                # Simulation horizon (s)
dt = 0.02               # Time step (s)
ts = 1:Int(Tmax/dt)     # Array of time indices

Now we can generate the training data. Suppose the box always starts at rest from the zero position, and the goal position can be anywhere in the range $q_\mathrm{ref} \in [-1,1]$. Our training data consists of a batch of 80 randomly-sampled goal positions and corresponding reference forces $u_\mathrm{ref}$.
[language = Julia]
nx, nref, batches = 2, 1, 80
x0 = zeros(nx, batches)
qref = 2*rand(nref, batches) .- 1
uref = k*qref

It is good practice (and faster) to simulate all simulation batches at once, so we define our dynamics functions to operate on matrices of states and controls. Each row corresponds to a different state or control, and each column corresponds to a simulation for a particular goal position.
[language = Julia]
f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] -
k*x[1:1,:] - μ*x[2:2,:].^2)/m]
fd(x::Matrix,u::Matrix) = x + dt*f(x,u)

RL problems typically involve simulating the system over some time horizon and collecting rewards or costs at each time step. Control policies are trained using approximations of the cost gradient $\nabla J_\theta$, as it is often difficult (or impossible) to compute the exact gradient due to the complexity of dynamics simulators. We refer the reader to  for further details <cit.>.

For this simple example, we can back-propagate directly through the dynamics function . The simulator below takes a batch of initial states, goal positions, and a controller  whose inputs are $[x; q_\mathrm{ref}]$. It computes trajectories of states and controls $z = \{[x_0;u_0], \ldots, [x_{T-1};u_{T-1}]\}$. To avoid the well-known issue of array mutation being unsupported in differentiable programming, we use a  to iteratively store the outputs <cit.>.

[language = Julia]
using Zygote: Buffer

function rollout(model, x0, qref)
z = Buffer([zero([x0;qref])], length(ts))
x = x0
for t in ts
u = model([x;qref])
z[t] = vcat(x,u)
x = fd(x,u)
return copy(z)

After computing these trajectories, we will need a function to evaluate the cost given some weightings $c_1,c_2,c_3$.
[language = Julia]
using Statistics

weights = [10,1,0.1]
function _cost(z, qref, uref)
Δz = z .- [qref; zero(qref); uref]
return mean(sum(weights .* Δz.^2; dims=1))
cost(z::AbstractVector, qref, uref) =
mean(_cost.(z, (qref,), (uref,)))

§.§.§ Define a model

We will train an LBDN controller with a Lipschitz bound of $\gamma = 20$. Its inputs are the state $x_t$ and goal position $q_\mathrm{ref}$, while its outputs are the control force $u_t$. We have chosen a model with two hidden layers each of 32 neurons just as an example. For details on how Lipschitz bounds can be useful in learning robust controllers, please see <cit.>.

[language = Julia]
using RobustNeuralNetworks

T  = Float64
γ  = 20                 # Lipschitz bound
nu = nx + nref          # Inputs (x and reference)
ny = 1                  # Outputs (control action)
nh = fill(32, 2)        # Hidden layers
model_ps = DenseLBDNParamsT(nu, nh, ny, γ)

§.§.§ Define a loss function

In constructing a loss function for this problem, we refer to Section <ref>. The  contain all information required to define a dense LBDN model. However,  is not a model that can be evaluated on data: it is a model parameterization, and contains the learnable parameters $\theta$. To train an LBDN given some data, we construct the model within the loss function using the  wrapper so that the mapping from direct to explicit parameters is captured during back-propagation. Our loss function therefore includes the following three components.

[language = Julia]
function loss(model_ps, x0, qref, uref)
model = LBDN(model_ps)            # Model
z = rollout(model, x0, qref)      # Simulation
return cost(z, qref, uref)        # Cost

§.§.§ Train the model

Having set up the RL problem, all that remains is to train the controller. The function below trains a model and keeps track of the training loss  (cost $J_\theta$) for each simulation in our batch of 80. Training is performed with the  optimizer over 250 epochs with a learning rate of $10^{-3}$.

[language = Julia]
using Flux

function train_box_ctrl!(
model_ps, loss_func;
epochs=250, lr=1e-3
costs = VectorFloat64()
opt_state = Flux.setup(Adam(lr), model_ps)
for k in 1:epochs

tloss, dJ = Flux.withgradient(
loss_func, model_ps, x0, qref, uref)
Flux.update!(opt_state, model_ps, dJ[1])
push!(costs, tloss)
return costs

costs = train_box_ctrl!(model_ps, loss)

§.§.§ Evaluate the trained model

We may now verify the performance of the trained model on a new set of reference positions. In the code below, we generate 60 batches of test data. In each one, the box starts at the origin at rest, and is moved through the fluid to a different (random) goal position $q_\mathrm{ref} \in [-1,1].$ We plot the states and controls alongside the loss curve from training in Figure <ref>. The box clearly moves to the required position within the time frame in all cases, experimentally verifying the performance of our controller.

model   = LBDN(model_ps)
x0_test = zeros(2,60)
qr_test = 2*rand(1, 60) .- 1
z_test  = rollout(model, x0_test, qr_test)

Loss curve and simulation results from the LBDN RL policy controlling the box system in Figure <ref>. The LBDN policy can push the box to any desired location in the domain of interest. The position and controller errors are $\Delta q$ and $\Delta u$ from Equation <ref>, respectively.

§.§.§ Advantages of separate parameters and models

As discussed in Section <ref>, there is a trade-off between convenience and performance in . The  and  wrappers exist to allow users to train robust models in a -like manner. These wrappers convert a model parameterization to its explicit form each time they are called, hence the user does not have to re-construct the model in the loss function.
[language = Julia]
loss2(model, x0, qref, uref) =
cost(rollout(model, x0, qref), qref, uref)

The cost is computation speed, particular in an RL context. Careful inspection of the  function shows that the  is evaluated many times within the loss function before the learnable parameters are updated with . As discussed in the Section <ref>, the major computational bottleneck in training RENs and LBDNs is the conversion from learnable (direct) parameters to an explicit model. Constructing the model only when the parameters are updated therefore saves considerably on computation time, particularly for large models.

For example, suppose we train single-hidden-layer LBDNs with $n = 2, 2^2, \ldots, 2^9$ neurons over 100 epochs on our box RL problem, and log the time taken to train each model when using both  and .

[language = Julia]
function lbdn_compute_times(n; epochs=100)

# Build model params and a model
lbdn_ps = DenseLBDNParamsT(nu, [n], ny, γ)
diff_lbdn = DiffLBDN(deepcopy(lbdn_ps))

# Time with LBDN vs DiffLBDN (respectively)
t_lbdn = @elapsed (
train_box_ctrl!(lbdn_ps, loss; epochs))
t_diff_lbdn = @elapsed (
train_box_ctrl!(diff_lbdn, loss2; epochs))
return [t_lbdn, t_diff_lbdn]

# Evaluate computation time
# Run it once first for just-in-time compiler
ns = 2 .^ (1:9)
lbdn_compute_times(2; epochs=1)
comp_times = reduce(hcat, lbdn_compute_times.(ns))

The results are plotted in Figure <ref>. Even for a single-layer LBDN with $2^9 = 512$ neurons, it is clear that using  takes an order of magnitude longer to train than only constructing the  model each time the  function is called. If we were training dynamic models with REN, the computational overhead of using  instead of  would be even more extreme, since the conversion from direct to explicit parameters in a REN is typically more computationally expensive than for LBDNs. It is for this reason that we strongly recommend using the  and  wrappers if many evaluations of the model are required before  (or equivalent) is called, as in RL.

Training time as a function of hidden-layer size for a single-hidden-layer LBDN constructed with both the  and  wrappers. Using the  wrapper for RL is significantly more efficient than re-constructing the explicit model at every evaluation of the  model.

§.§ Observer design

In Section <ref>, we designed a controller for a simple nonlinear system assuming that the controller had full state knowledge: that is, it had access to both the position and velocity of the box. In many practical situations, we may only be able to measure some of the system states. For example, our box may have a camera to estimate its position but not its velocity. In these cases, we need a state observer to estimate the full state of the system for feedback control.

In this example, we will show how a contracting REN can be used to learn stable observers for dynamical systems. A common approach to designing state estimators for nonlinear systems is the Extended Kalman Filter (EKF). In our case, we will consider observer design as a supervised learning problem. For a detailed explanation of the theory behind learning state observers, and for a similar example designing an observer for a Partial Differential Equations (PDE), please refer to Section VIII of <cit.>.

§.§.§ Background theory

We briefly summarise some background theory from <cit.> relevant to this example. Suppose we have a discrete-time, nonlinear dynamical system of the form

\begin{align}
x_{t+1} &= f_d(x_t, u_t) \\
y_t &= g_d(x_t, u_t)
\end{align}

with state vector $x_t,$ controlled inputs $u_t,$ and measured outputs $y_t.$ Our aim is to estimate the sequence $\{x_0, x_1, \ldots, x_T \}$ over some time period $[0,T]$ given only the measurements $y_t$ and inputs $u_t$ at each time step. We will use a very general form for an observer

\begin{equation}
\hat{x}_{t+1} = f_o(\hat{x}_t, u_t, y_t)
\end{equation}

where $\hat{x}_t$ is the state estimate. A more common (but more restrictive) structure is the well-known Luenberger observer <cit.>.

To estimate the true state, our observer error $(x_t - \hat{x}_t)$ must converge to zero as time progresses, or $\hat{x}_t \rightarrow x_t$ as $t \rightarrow \infty$. As outlined in <cit.>, our observer only has to satisfy the following two conditions to guarantee this.

* The observer must be a contracting system (Sec. <ref>).

* The observer must satisfy a “correctness” condition which says that, given perfect knowledge of the state, measurements, and inputs, the observer can exactly predict the next state. Mathematically, we write this as

\begin{equation}
f_o(x_t,u_t,y_t) = f_d(x_t,u_t)
\end{equation}

where $y_t = g_d(x_t,u_t)$. Note the use of $x_t$ not $\hat{x}_t$. It turns out that if the correctness condition is only approximately satisfied such that $|f_o(x_t,u_t,y_t) - f_d(x_t,u_t)| < \rho$ for some small number $\rho\in\mathbb{R}$, then the observer error will still be bounded. See Appendix E of <cit.> for details.

The first condition, contraction, is already guaranteed for all REN models in . Therefore, to learn a stable observer with RENs, our only requirement is to minimise the one-step-ahead prediction error to approximate the correctness condition. If we have a batch of data $z = \{x_i, u_i, y_i, \ i = 1,2,\ldots,N\},$ this corresponds to minimising the loss function

\begin{equation}
\mathcal{L}(z, \theta) = \sum_{i=1}^N |f_o(x_i,u_i,y_i) - f_d(x_i,u_i)|^2,
\end{equation}

where $\theta$ contains the learnable parameters of the REN.

§.§.§ Generate training data

Consider the same nonlinear box system from Section <ref>, with the a change in setup so that we can only measure the box position. We introduce a measurement function  such that $y_t = x_t$.

[language = Julia]
m = 1                   # Mass (kg)
k = 5                   # Spring constant (N/m)
μ = 0.5                 # Viscous damping (kg/m)
nx = 2                  # Number of states

# Continuous and discrete dynamics and measurements
f(x::Matrix,u::Matrix) = [x[2:2,:]; (u[1:1,:] -
k*x[1:1,:] - μ*x[2:2,:].^2)/m]
fd(x,u) = x + dt*f(x,u)
gd(x::Matrix) = x[1:1,:]

For this example, we assume that the box always starts at rest in a random initial position between $\pm0.5$m, after which it is released and allowed to oscillate freely with no added forces (so $u = 0$). Learning an observer typically requires a large amount of training data to fully capture the behaviour of the system, hence we consider 200 batches each simulating 10 s of motion.
[language = Julia]
Tmax = 10               # Simulation horizon
dt = 0.01               # Time-step (s)
ts = 1:Int(Tmax/dt)     # Time array indices

# Generate batches of training data
nbatch = 200
u = fill(zeros(1, nbatch), length(ts)-1)
X = fill(zeros(1, nbatch), length(ts))
X[1] = 0.5*(2*rand(nx, nbatch) .- 1)

for t in ts[1:end-1]
X[t+1] = fd(X[t],u[t])

We have stored the states of the system across each batch in . To compute the one-step-ahead loss $\mathcal{L},$ we will need to separate this data into the states at the “current” time step  and at the “next” time step , then compute the measurement outputs. We then store the data for training, shuffling it so there is no bias in the training towards earlier time steps.
[language = Julia]
using Random

# Current/next state, measurements
Xt = X[1:end-1]
Xn = X[2:end]
y  = gd.(Xt)

# Store training data
obsv_data = [[ut; yt] for (ut,yt) in zip(u, y)]
indx = shuffle(1:length(obsv_data))
data = zip(Xn[indx], Xt[indx], obsv_data[indx])

§.§.§ Define a model

We can construct the parameterization for a contracting REN model using . The inputs to the model are $[u_t;y_t]$, and its outputs are the next state estimate $\hat{x}_{t+1}$. The flag  sets the output map of the REN to just return its own internal state – i.e: $C_2 = I$, $D_{21} = 0$, $D_{22} = 0$, $b_y = 0$ from Equation <ref>. This makes the internal state of the REN exactly the state estimate $\hat{x}_t$.

[language = Julia]
using RobustNeuralNetworks

T  = Float64
nv = 100
nu = size(obsv_data[1], 1)
ny = nx
model_ps = ContractingRENParamsT(
nu, nx, nv, ny; output_map=false)
model = DiffREN(model_ps)

§.§.§ Define a loss function

As outlined in Section <ref>, our loss function should be the one-step-ahead prediction error of the REN observer. We write this as follows, noting that all subtypes of  return both their updated internal state and their output (in that order).
[language = Julia]
using Statistics

function loss(model, xn, xt, inputs)
xpred = model(xt, inputs)[1]
return mean(sum((xn - xpred).^2, dims=1))

§.§.§ Train the model

The function below trains the observer with the  optimizer over 100 epochs and decreases the maximum learning rate from $10^{-3}$ to $10^{-4}$ if the mean loss stops decreasing between epochs. The core of this function is a simple  training loop, expanded out for clarity.
[language = Julia]
using Flux

function train_observer!(
model, data;
epochs=100, lr=1e-3, min_lr=1e-4
opt_state = Flux.setup(Adam(lr), model)
mean_loss = [1e5]
for epoch in 1:epochs

# Gradient descent update
batch_loss = []
for (xn, xt, inputs) in data
tloss, dJ = Flux.withgradient(
loss, model, xn, xt, inputs)
Flux.update!(opt_state, model, dJ[1])
push!(batch_loss, tloss)

# Reduce lr if loss is stuck or growing
push!(mean_loss, mean(batch_loss))
if (mean_loss[end] >= mean_loss[end-1])
(lr > min_lr)
lr *= 0.1
Flux.adjust!(opt_state, lr)
return mean_loss
tloss = train_observer!(model, data)

§.§.§ Evaluate the trained model

We have trained the REN observer to minimise the one-step-ahead prediction error, but we are yet to test whether the the observer error actually does converge to zero. We set up the following 50 batches of test data as a demonstration.
[language = Julia]
nbatch    = 50
ts_test   = 1:Int(10/dt)
u_test    = fill(zeros(1, nbatch), length(ts_test))
x_test    = fill(zeros(nx,nbatch), length(ts_test))
x_test[1] = 0.2*(2*rand(nx, nbatch) .-1)

for t in ts_test[1:end-1]
x_test[t+1] = fd(x_test[t], u_test[t])
y_test = gd.(x_test)
obsv_in = [[u;y] for (u,y) in zip(u_test, y_test)]

Next, we need a function to simulate the REN observer using its own state $\hat{x}_t$ rather than the true system state $x_t$, which was used for training. We use the very neat tool  for this. We assume that the observer has no knowledge of the initial state and simply guesses $\hat{x}_0 = 0$ for all 50 batches.
[language = Julia]
function simulate(model::AbstractREN, x0, u)
recurrent = Flux.Recur(model, x0)
output = recurrent.(u)
return output
x0hat = zeros(model.nx, nbatch)
xhat = simulate(model, x0hat, obsv_in)

Simulation results showing the observer predictions and observer error with the box starting at 50 different initial conditions. The left panels compare the true (grey) and estimated (red) states, while the right panels show the observer error $x - \hat{x}$ over time. The observer error converges for all 50 test cases.

The results are plotted in Figure <ref>. In the left-hand panels, the observer predictions (red) almost exactly match the true states (grey) after approximately 4 s. This is confirmed by the right-hand panels, which show the observer error $x_t - \hat{x}_t$ smoothly converging to zero as the observer estimates the correct states for all simulations.

It is worth noting that at no point did we directly train the REN to minimise the observer error. This is a natural result of using a model that is guaranteed to be contracting, and training it to minimise the one-step-ahead prediction error. There is still some residual observer error in the velocity in Figure <ref>, since our observer was only trained to approximately satisfy the correctness condition. However, this could easily be reduced or eliminated using a larger observer model, more training data, and more training epochs.

§ SUMMARY AND CONCLUSIONS

This paper has presented , a Julia package for robust machine learning based on the recently-proposed Recurrent Equilibrium Network (REN) and Lipschitz-Bounded Deep Network (LBDN) model classes. The models are unique in that they naturally satisfy a set of built-in robustness certificates, such as contraction and Lipschitz bounds. We have presented an overview of the model architectures, including background theory on robustness metrics in nonlinear systems, and have outlined the package structure and its usage alongside Julia's main machine-learning library, . We have demonstrated via examples in image classification, reinforcement learning, and observer design that the package is easy to use in many common machine learning problems, while also offering the advantage of robustness guarantees.

We intend  to be widely-applicable in the scientific and machine learning communities for learning-based problems in which robustness certificates are crucial, and have already used the package in our own research in robust reinforcement learning <cit.>. Some areas in which this package will be most applicable include: data-driven control and state estimation, image classification and segmentation, and privacy and security. We intend to expand the package with more robust neural network architectures in the future. Examples include LBDNs with one-dimensional convolution <cit.> and circular convolutions <cit.>, continuous-time REN models <cit.>, and RENs respecting other non-Euclidean contraction metrics <cit.>. We encourage any and all contributions to  to further its use in robust machine learning problems.

This work was supported in part by the Australian Research Council (DP190102963).