Neural networks, for example image classifiers, are often vulnerable to attacks where small, sub-perceptual changes to the input cause misclassification.

Informally, given an input $x$ to a neural network $f$, the verification problem is to prove that for any $x'$ that is a bounded perturbation of $x$, $f(x') = f(x)$. Successful proofs can rule out attacks like the one above.

The current state of the art is an algorithm called $\alpha\beta$-CROWN, an implementation of which has won the International Verification of Neural Networks Competition for the past five years. I think $\alpha\beta$-CROWN is a lovely algorithm in an important research area where progress is unlikely in my lifetime.

I implemented $\alpha\beta$-CROWN for a deep-learning library called tinygrad and this blog post is about how it works. This is for the most part a synthesis of these three papers, my very modest contribution being a description of how the algorithm handles shape changes (reshape/flatten/permute) which does not appear in existing descriptions.

Background

$x'$ is a bounded perturbation of $x$ if it is an element of an $\ell_p$-ball of radius $\epsilon$ centered at $x$, denoted $B_p(x, \epsilon)$:

where $\lVert x \rVert_p$ is $x$'s $p$-norm,

A convenient fact is that there is a closed-form solution to the minimum of any linear (or affine) function over an $\ell_p$ ball (Hölder's inequality):

where $\langle x, y\rangle$ denotes the dot product of vectors $x$ and $y$, and $\lVert\cdot\rVert_q$ is the dual norm of $\lVert\cdot\rVert_p$, defined by $\frac{1}{q} + \frac{1}{p} = 1.$

In image classification, for an image $x$, a neural network typically outputs an $n$-dimensional vector $f(x)\in\mathbb{R}^n$ and $x$'s classification is said to be the index in $f(x)$ with the largest value. Thus, for the example in Figure 1, proving no bounded perturbation of an image $x$ gets classified as 3 is equivalent to checking that $\forall x'\in B_p(x, \epsilon), f(x')_7-f(x')_3 > 0.$ In turn, a common approach to neural network verification (and the one taken by $\alpha\beta$-CROWN) is to compute affine bounds on the output of $f$, then use Hölder's inequality to minimize $f(x')_7-f(x')_3$ over $B_p(x, \epsilon)$. If the minimum is $>0$, then no images in a bounded perturbation of $x$ will be classified as 3.

$\alpha\beta$-CROWN on Graphs

Many deep-learning libraries represent neural networks as graphs where edges correspond to data and vertices to differentiable operations. For example, $f(x) = h(x) + g(x)$ might be represented as

To implement $\alpha$-CROWN on a graph, we need a way to compute bounds on the output of each operator given bounds on its input, then bounds on the output of the network can be propagated through the graph and written in terms of the network's input. For example, for the + operator above, if

are affine lower bounds on $h$ and $g$, then

is a lower bound on $h(x) + g(x)$.

Handling ReLU

Nonlinear functions like ReLU are more challenging to bound. $\alpha\beta$-CROWN uses linear over/under-approximations.

Figure 3. The dashed lines are linear bounds on a ReLU over the interval $[\ell, u]$.

First, Hölder's inequality is used to compute concrete bounds $\ell$ and $u$ on the ReLU's input. Then, when $\ell < 0 < u$ (i.e. the ReLU is nonlinear), for some $0 \le \alpha$, letting $\sigma$ denote ReLU, the following equations bound it:

Notice how any value of $\alpha$ between zero and one gives a lower bound when $\ell < 0 < u$. Thus, towards choosing a good value for $\alpha$, notice that the ReLU bounds, the $+$ function bounds, and Hölder's inequality are all computed using differentiable operations. In fact, in general, all the equations for bound propagation in $\alpha\beta$-CROWN are differentiable. In turn, given bounds on a neural network's output in terms of different ReLU's $\alpha$ values, we can use gradient descent to find $\alpha$ values that tighten them. I think this is cool!

$\beta$-CROWN

In Figure 3, if $\ell \ge 0$ or $u \le 0$, then setting $\alpha$ to 1 or 0 makes the ReLU bounds exact; ReLU's with this property are called stable. If we can compute bounds on $f$ conditional on a ReLU being stable, we could consider both cases and merge them. For example, if $\ell^{(+)},u^{(+)}$ were lower and upper bounds on a ReLU when $\ell \ge 0$, and $\ell^{(-)},u^{(-)}$ were bounds when $u \le 0$, then

are overall bounds on the ReLU. This is the idea behind $\beta$-CROWN. To propagate bounds on $f^{L-1}(x)$ to $f^L(x)$,

subject to constraints $C$, we let $S f^{L-1}(x)$ be positive if $f^{L-1}(x)$ satisfies $C$ and negative otherwise. For example, in the $\ell^{(+)}$ case where the constraint enforces $f^{L-1}(x)\ge 0$, we have $S=1$. Then, if

by Lagrangian relaxation and weak duality,

where $\beta \ge 0$. To find a value for $\beta$ that gives a tight bound we can, once again, use gradient descent. Notice that the $f^{L-1}(x)$ term can be removed by recursively getting bounds on $\beta S f^{L-1}(x)$ in terms of $x$. Bounds for the $(+)$ and $(-)$ cases can be merged to get overall bounds on the ReLU as described above. Finding upper bounds is a variant of the same procedure as $\min f(x)=-\max -f(x)$.

tinygrad

tinygrad is a popular deep-learning library with lazy evaluation and a small set of UOps to which tensor operations are lowered before being run. Because tinygrad is lazy, tensors are represented as a graph that computes their value on demand, making accessing a computation's graph straightforward. Furthermore, tinygrad's small set of UOps means bound propagation only needs to be implemented for this small operation set to support complex networks. These properties make tinygrad amenable to an implementation of $\alpha\beta$-CROWN.

However, tinygrad's UOps are distinct from the operations described in previous work on $\alpha\beta$-CROWN. For example, whereas $\alpha$-CROWN considers matrix multiply a primitive operation, tinygrad compiles matrix multiply to a series of low-level shape changes. Specifically, for $A,B\in\mathbb{R}^{n\times n}$, tinygrad computes $C=AB$ via:

While available implementation of $\alpha\beta$-CROWN handle shape changes, published descriptions of the algorithm don't, instead describing how to compute bounds on $f(x)$ in terms of vectors $x, \underline{b}, \overline{b}$ and matrices $\underline{W}, \overline{W}$:

When the input to the network is not a vector (for example, an image with shape $(\texttt{bs, w, h, ch})$), the matrix multiply above is no longer semantically valid. Thus, beyond an $\alpha\beta$-CROWN implementation for tinygrad, this blog's small contribution in the following section is a description of how to compute $\alpha\beta$-CROWN bounds for arbitrarily shaped inputs subject to arbitrary shape changes. The method is semantically identical to the one used by $\alpha\beta$-CROWN's existing PyTorch implementation.

Handling Shape Changes

To handle high-rank tensors and shape changes in tinygrad, bounds on the network are written in terms of tensor contractions.

For $x$ with shape $(a_1,\dots,a_A)$ and $T$ with shape $(p_1,\dots,p_P,a_1,\dots,a_A)$ the tensor contraction of $T$ and $x$ is written,

Intuitively, $T_{i_1,\ldots,i_P}$ can be thought of as describing the linear combination of elements of $x$ that make up index $i_1,\ldots,i_P$ in the result.

To handle shape changes like permute, expand, reshape, and shrink notice that all of these can be described by a function $\pi$ mapping an index $i$ in the resulting tensor to an index $\pi(i)$ in the original tensor. Applying a shape change to the bounds is thus equivalent to moving the linear combination of elements of $x$ used to compute index $\pi(i)$ to index $i$, i.e. applying the shape change to $T$'s leading $p_1,\dots,p_P$ dimensions. I think this is quite elegant and was happy to figure it out.

The code is available here.

Thanks to Dhruv for working on this with me, Michael Everett for a great class, and Nick and Noah for interesting conversations.