Cryptanalytic Extraction of Neural Network Models

Introduction

The past decade has seen significant advances in machine learning, and deep learning in particular. Tasks viewed as being completely infeasible at the beginning of the decade became almost completely solved by the end. AlphaGo [SHM+16] defeated professional players at Go, a feat in 2014 seen as being at least ten years away [Lev14]. Accuracy on the ImageNet recognition benchmark improved from $73\%$ in 2010 to $98.7\%$ in 2019, a $20\times$ reduction in error rate [XLHL20]. Neural networks can generate photo-realistic high-resolution images that humans find indistinguishable from actual photographs [KLA+19]. Neural networks achieve higher accuracy than human doctors in limited settings, such as early cancer detection [EKN+17].

These advances have brought neural networks into production systems. The automatic speech recognition systems on Google’s Assistant, Apple’s Siri, and Amazon’s Alexa are all powered by speech recognition neural networks. Neural Machine Translation [BCB15] is now the technique of choice for production language translation systems [WSC+16]. Autonomous driving is only feasible because of these improved image recognition neural networks.

High-accuracy neural networks are often held secret for at least two reasons. First, they are seen as a competitive advantage and are treated as trade secrets [Wen90]; for example, none of the earlier systems are open-source. Second, is seen as improving both security and privacy to keep these models secret. With full white-box access it is easy to mount evasion attacks and generate adversarial examples [SZS+14, BCM+13] against, for instance, abuse- or spam-detection models. Further, white-box access allows model inversion attacks [FJR15]: it is possible to reconstruct identifiable images of specific people given a model trained to recognize specific human faces. Similarly, given a language model trained on text containing sensitive data (e.g., credit card numbers), a white-box attacker can pull this sensitive data out of the trained model [CLE+19].

Fortunately for providers of machine learning models, it is often expensive to reproduce a neural network. There are three reasons for this cost: first, most machine learning requires extensive training data that can be expensive to collect; second, neural networks typically need hyper-parameter tuning requiring training many models to identify the optimal final model configuration; and third, even performing a final training run given the collected training data and correctly configured model is expensive.

For all of the above reasons, it becomes clear that (a) adversaries are motivated for various reasons to obtain a copy of existing deployed neural networks, and (b) preserving the secrecy of models is highly important. In practice companies ensure the secrecy of these models by either releasing only an API allowing query access, or releasing on-device models, but attempting to tamper-proof and obfuscate the source to make it difficult to extract out of the device.

Understandably, the above weak forms of protection are often seen as insufficient. The area of “secure inference” improves on this by bringing tools from Secure Function Evaluation (SFE), which allows mutually distrustful cooperating parties to evaluate $f(x)$ where $f$ is held by one party and $x$ by the other. The various proposals often apply fully homomorphic encryption [Gen09, GBDL+16], garbled circuits [Yao86, RWT+18], or combinations of the two [MLS+20]. Per the standard SFE guarantee, secure inference “does not hide information [about the function $f$] that is revealed by the result of the prediction” [MLS+20]. However this line of work often implicitly assumes that total leakage from the predictions is small, and that recovering the function from its output would be difficult.

In total, it is clear that protecting the secrecy of neural network models is seen as important both in practice and in the academic research community. This leads to the question that we study in this paper:

Is it possible to extract an identical copy of a neural network given oracle (black-box) access to the target model?

While this question is not new [TZJ+16, MSDH19, JCB+20, RK20], we argue that model extraction should be studied as a cryptanalytic problem. To do this, we focus on model extraction in an idealized environment where a machine learning model is made available as an oracle $\mathcal{O}$ that can be queried, but with no timing or other side channels. This setting captures that of obfuscated models made public, prediction APIs, and secure inference.

The key insight of this paper is that model extraction is closely related to an extremely well-studied problem in cryptography: the cryptanalysis of blockciphers. Informally, a symmetric-key encryption algorithm is a keyed function $E_{k}\colon\mathcal{X}\to\mathcal{Y}$ that maps inputs (plaintexts) $x\in\mathcal{X}$ to outputs (ciphertexts) $y\in\mathcal{Y}$. We expect all practically important ciphers to be resistant, at the very least, to key recovery under the adaptive chosen-plaintext attack, i.e., given some bounded number of (adaptively chosen) plaintext/ciphertext pairs $\{(x_{i},y_{i})\}$ an encryption algorithm is designed so that the key $k$ cannot be extracted by a computationally-bounded adversary.

Contrast this to machine learning. A neural network model is (informally) a parameterized function $f_{\theta}\colon\mathcal{X}\to\mathcal{Y}$ that maps input (e.g., images) $x\in\mathcal{X}$ to outputs (e.g., labels) $y\in\mathcal{Y}$. A model extraction attack adaptively queries the neural network to obtain a set of input/output pairs $\{(x_{i},y_{i})\}$ that reveals information about the weights $\theta$. Neural networks are not constructed by design to be resistant to such attacks.

Thus, viewed appropriately, performing a model extraction attack—learning the weights $\theta$ given oracle access to the function $f_{\theta}$—is a similar problem to performing a chosen-plaintext attack on a nontraditional “encryption” algorithm.

Given that it took the field of cryptography decades to design encryption algorithms secure against chosen-plaintext attacks, it would be deeply surprising if neural networks, where such attacks are not even considered in their design, were not vulnerable. Worse, the primary objective of cipher design is robustness against such attacks. Machine learning models, on the other hand, are primarily designed to be accurate at some underlying task, making the design of chosen-plaintext secure neural networks an even more challenging problem.

There are three differences separating model extraction from standard cryptanalysis that make model extraction nontrivial and interesting to study.

First, the attack success criterion differs. While a cryptographic break can be successful even without learning key bits—for example by distinguishing the algorithm from a pseudo-random function, only “total breaks” that reveal (some of) the actual model parameters $\theta$ are interesting for model extraction.

Second, the earlier analogy to keyed ciphers is imperfect. Neural networks typically take high-dimensional inputs (e.g., images) and return low-dimensional outputs (e.g., a single probability). It is almost more appropriate to make an analogy to cryptanalysis of keyed many-to-one functions, such as MACs. However, the security properties of MACs are quite different from those of machine learning models, for example second preimages are expected rather than shunned in neural networks.

Finally, and the largest difference in practice, is that machine learning models deal in fixed- or floating-point reals rather than finite field arithmetic. As such, there are many components to our attack that would be significantly simplified given infinitely precise floating-point math, but given the realities of modern machine learning, require far more sophisticated attack techniques.

2 Our Results

We introduce a differential attack that is effective at performing functionally-equivalent neural network model extraction attacks. Our attack traces the neural network’s evaluation on pairs of examples that differ in a few entries and uses this to recover the layers (analogous to the rounds of a block cipher) of a neural network one by one. To evaluate the efficacy of our attack, we formalize the definition of fidelity introduced in prior work [JCB+20] and quantify the degree to which a model extraction attack has succeeded:

Two models $f$ and $g$ are $(\epsilon,\delta)$-functionally equivalent on $S$ if

Table 1 reports the results of our differential attack across a wide range of model sizes and architectures, reporting both $(\varepsilon,\delta)$-functional equivalence on the set $S=^{d_{0}}$, the input space of the model, along with a direct measurement of $\max\,\,\lvert\theta-\hat{\theta}\rvert$, directly measuring the error between the actual model weights $\theta$ and the extracted weights $\hat{\theta}$ (as described in Section 6.2).

The remainder of this paper is structured as follows. We introduce the notation, threat model, and attacker goals and assumptions used in Section 2. In Section 4 we introduce an idealized attack that extracts $(0,0)$-functionally-equivalent neural networks assuming infinite precision arithmetic. Section 5 develops an instantiation of this attack that works in practice with finite-precision arithmetic to yield $(\varepsilon,\delta)$-functionally equivalent attacks.

3 Related Work

Model extraction attacks are classified into two categories [JCB+20]: task accuracy extraction and fidelity extraction. The first paper to study task accuracy extraction [TZJ+16] introduced techniques to steal similar models that approximately solve the same underlying decision task on the natural data distribution, but do not necessarily match the predictions of the oracle precisely. While further work exists in this space [CCG+20, KTP+20], we instead focus on fidelity extraction where the adversary aims to faithfully reproduce the predictions of the oracle model, when it is incorrect with respect to the ground truth. Again, [TZJ+16] studied this problem and developed (what we would now call) functionally equivalent extraction for the case of completely linear models.

This attack was then extended by a theoretical result defining and giving a method for performing functionally-equivalent extraction for neural networks with one layer, assuming oracle access to the gradients [MSDH19]. A concrete implementation of this one layer attack that works in practice, handling floating point imprecision, was subsequently developed through applying finite differences to estimate the gradient [JCB+20]. Parallel work to this also extended on these results, focusing on deeper networks, but required tens to hundreds of millions of queries [RK20]; while the theoretical results extended to deep networks, the implementation in practice only extracts up to the first two layers. Our work builds on all of these four results to develop an approach that is $10^{6}$ times more accurate, requiring $10^{3}$ times fewer queries, and applies to larger models.

Even without query access, it is possible to steal models with just a cache side-channel [BBJP19], although with less fidelity than our attack that we introduce which are $2^{20}\times$ more precise. Other attacks target hyperparameter extraction—that is, extracting high-level details about the model: through what method it was trained, if it contains convolutions, or related questions [WG18]. It is further possible to steal hyperparameters with cache side channels [HDK+20].

Recent work has studied the learnability of deep neural networks with random weights in the statistical query (SQ) model [DGKP20], showing that learnability drops off exponentially with the depth of the network. This line of work does not address the cryptographic hardness of extraction in the non-SQ model—precisely the question addressed in this work in the empirical setting.

While not directly related to our problem, it is worth noting that we are not the first to treat neural networks as just another type of mathematical function that can be analyzed without any specific knowledge of machine learning. Shamir et al. [SSRD19] explain the existence of adversarial examples [SZS+14, BCM+13], which capture evasion attacks on machine learning classifiers, by considering an abstract model of neural networks.

In a number of places, our attack draws inspiration from the cryptanalysis of keyed block-ciphers, most prominently differential cryptanalysis [BS91]. We neither assume nor require familiarity with this field, but the informed reader may enjoy certain parallels.

Preliminaries

This paper studies an abstraction of neural networks as functions $f\colon\mathcal{X}\to\mathcal{Y}$. Our results are independent of any methods for selecting the function $f$ (e.g., stochastic gradient descent), and are independent of any utility of the function $f$. As such, machine learning knowledge is neither expected nor necessary.

A $k$-deep neural network $f_{\theta}(x)$ is a function parameterized by $\theta$ that takes inputs from an input space $\mathcal{X}$ and returns values in an output space $\mathcal{Y}$. The function $f$ is composed as a sequence of functions alternating between linear layers $f_{j}$ and a nonlinear function (acting component-wise) $\sigma$:

While representing each layer $f_{j}$ as a full matrix product is the most general definition of a layer, which is called fully connected, often layers have more structure. For example, it is common to use (discrete) convolutions in neural networks that operate on images. Convolutional layers take the input as a $n\times m$ matrix and convolve it with a kernel, such as a $3\times 3$ matrix. Importantly, however, it is always possible to represent a convolution as a matrix product.

The neurons $\{\eta_{i}\}_{i=1}^{N}$ are functions receiving an input and passing it through the activation function $\sigma$. There are a total of $N=\sum_{j=1}^{k-1}d_{j}$ neurons.

In this paper we exclusively study the ReLU [NH10] activation function, given by $\sigma(x)=\max(x,0)$. Our results are a fundamental consequence of the fact that ReLU neural networks are piecewise linear functions.

The architecture of a neural network captures the structure of $f$: (a) the number of layers, (b) the dimensions of each layer $\{d_{i}\}_{i=0}^{k}$, and (c) any additional constraints imposed on the weights $A^{(i)}$ and biases $b^{(i)}$.

We use the shorthand $a$-$b$-$c$ neural network to denote the sizes of each dimension; for example a 10-20-5 neural network has input dimension 10, one layer with 20 neurons, and output dimension 5. This description completely characterizes the structure of $f$ for fully connected networks. In practice, there are only a few architectures that represent most of the deployed deep learning models [ZL17], and developing new architectures is an extremely difficult and active area in research [HZRS16, SIVA17, TL19].

The parameters $\theta$ of $f_{\theta}$ are the concrete assignments to the weights $A^{(j)}$ and biases $b^{(j)}$, obtained during the process of training the neural network.

It is beyond the scope of this paper to describe the training process which produces the parameters $\theta$: it suffices to know that the process of training is often computationally expensive and that training is a nondeterministic process, and so training the same model multiple times will give different sets of parameters.

2 Adversarial Goals and Resources

There are two parties in a model extraction attack: the oracle $\mathcal{O}$ who returns $f_{\theta}(x)$, and the adversary who generates queries $x$ to the oracle.

A model parameter extraction attack receives oracle access to a parameterized function $f_{\theta}$ (in our case a $k$-deep neural network) and the architecture of $f$, and returns a set of parameters $\hat{\theta}$ with the goal that $f_{\theta}(x)$ is as similar as possible to $f_{\hat{\theta}}(x)$.

Throughout this paper we use the $\hat{\_}$ symbol to indicate an extracted parameter. For example, $\hat{\theta}$ refers to the extracted weights of a model $\theta$.

There is a spectrum of similarity definitions between the extracted weights and the oracle model that prior work has studied [TZJ+16, JCB+20, KTP+20]; we focus on the setting where the adversarial advantage is defined by $(\varepsilon,\delta)$-functionally equivalent extraction as in Definition 1.

Analogous to cryptanalysis of symmetric-key primitives, the degree to which a model extraction attack succeeds is determined by (a) the number of chosen inputs to the model, and (b) the amount of compute required.

We make several assumptions of the oracle $\mathcal{O}$ and the attacker’s knowledge. (We believe many of these assumptions are not fundamental and can be relaxed. Removing these assumptions is left to future work.)

Architecture knowledge. We require knowledge of the architecture of the neural network.

Full-domain inputs. We feed arbitrary inputs from $\mathcal{X}$.

Complete outputs. We receive outputs directly from the model $f$ without further processing (e.g., by returning only the most likely class without a score).

Precise computations. $f$ is specified and evaluated using 64-bit floating-point arithmetic.

ReLU Activations. All activation functions ($\sigma$’s) are ReLU’s.This is the only assumption fundamental to our work. Switching to any activation that is not piecewise linear would prevent our attack. However, as mentioned, all state-of-the-art models use exclusively (piecewise linear generalizations of) the ReLU activation function [SIVA17, TL19].

Overview of the Differential Attack

Given oracle access to the function $f_{\theta}$, we can estimate $\partial f_{\theta}$ through finite differences along arbitrary directions. For simple linear functions defined by $f(x)=a\cdot x+b$, its directional derivative satisfies $\frac{\partial f}{\partial e_{i}}\equiv a_{i}$, where $e_{i}$ is the basis vector and $a_{i}$ is the $i$th entry of the vector $a$, allowing direct recovery of its weights through querying on these well-chosen inputs.

In the case of deep neural networks, we consider second partial directional derivatives. ReLU neural networks are piecewise linear functions with ${\partial^{2}f\over\partial x^{2}}\equiv 0$ almost everywhere, except when the function has some neuron $\eta_{j}$ at the boundary between the negative and positive region (i.e., is at its critical point). We show that the value of the partial derivative $\partial^{2}f\over\partial e_{i}^{2}$ evaluated at a point $x$ so that neuron $\eta_{j}$ is at such a critical point actually directly reveals the weight $T(A^{(1)}_{i,j})$ for some transform $T$ that is invertible—and therefore the adversary can learn $A^{(1)}_{i,j}$. By repeating this attack along all basis vectors $e_{i}$ and for all neurons $\eta_{j}$ we can recover the complete matrix $A^{(1)}$. Once we have extracted the first layer’s weights, we are able to “peel off” that layer and re-mount our attack on the second layer of the neural network, repeating to the final layer. There are three core technical difficulties to our attack:

For each neuron $\eta$, our attack does not exactly recover $A^{(l)}_{i}$, the $i$th row of $A^{(l)}$, but instead a scalar multiple $v=\alpha\cdot A^{(l)}_{i}$. While losing a constant $\alpha>0$ keeps the neural network in the same equivalence class, the sign of $\alpha$ is important and we must distinguish between the weight vector $A^{(l)}_{i}$ and $-A^{(l)}_{i}$. We construct two approaches that solve this problem, but in the general case we require exponential work (but a linear number of queries).

On the first layer, we can directly compute the derivative entry-by-entry, measuring ${\partial^{2}f\over\partial e_{i}^{2}}$ for each standard basis vector $e_{i}$ in order to recover $A^{(1)}_{ij}$. Deeper in the network, we can not move along standard basis vector vectors. Worse, for each input $x$ on average half of the neurons are in the negative region and thus their output is identically ; when this happens it is not possible to learn the weight along edges with value zero. Thus we are required to develop techniques to elicit behavior from every neuron, and techniques to cluster together partial recoveries of each row of $A^{(l)}_{i}$ to form a complete recovery.

Implementing our attack in practice with finite precision neural networks introduces additional complexity. In order to estimate the second partial derivative, we require querying on inputs that differ by only a small amount, reducing the precision of the extracted first weight matrix to twenty bits, or roughly $10^{-6}$. This error of $10^{-6}$ is not large to begin with, but this error impacts our ability to recover the next layer, compounding multiplicatively the deeper we go in the network. Already in the second layer, the error is magnified to $10^{-4}$, which can completely prevent reconstruction for the third layer: our predicted view of the hidden state is sufficiently different from the actual hidden state that our attack fails completely. We resolve this through two means. First, we introduce numerically stable methods assuming that all prior layers have been extracted to high precision. Second, we develop a precision-refinement technique that takes a prefix of the first $j\leq k$ layers of a neural network extracted to $n$ bits of precision and returns the $j$-deep model extracted to $2n$ bits of precision (up to floating-point tolerance).

Idealized Differential Extraction Attack

We now introduce our $(0,0)$-functionally-equivalent model extraction attack that assumes infinite precision arithmetic and recovers completely functionally equivalent models. Recall our attack assumptions (Section 2.2); using these, we present our attack beginning with two “reduced-round” attacks on 0-deep (Section 4.1) and 1-deep (Section 4.2) neural networks, and then proceeding to $k$-deep extraction for contractive (Section 4.3) and expansive (Section 4.4) neural networks. Section 5 refines this idealized attack to work with finite precision.

Zero-deep neural networks are linear functions $f(x)\equiv A^{(1)}\cdot x+b^{(1)}$. Querying $d_{0}$ linearly independent suffices to extract $f$ by solving the resulting linear system.

However let us view this problem differently, to illuminate our attack strategy for deeper networks. Consider the parallel evaluations $f(x)$ and $f(x+\mathbf{\delta})$, with

This allows us to directly read off the weights of $A^{(1)}$. Put differently, we perform finite differences to estimate the gradient of $f$, given by $\nabla_{x}f(x)\equiv A^{(1)}$.

2 One-Deep Neural Network Extraction

Many of the important problems that complicate deep neural network extraction begin to arise at 1-deep neural networks. Because the function is no longer completely linear, we require multiple phases to recover the network completely. To do so, we will proceed layer-by-layer, extracting the first layer, and then use the 0-deep neural network attack to recover the second layer.

For the remainder of this paper it will be useful to have two distinct mental models of the problem at hand. First is the symbolic view shown previously in Figure 1. This view directly studies the flow of information through the neural networks, represented as an alternating sequence of linear layers and non-linear transformations. This view helps understanding the algebraic steps of our attack.

The second is the geometric view. Because neural networks operate over the real vector space, they can be visualized by plotting two dimensional slices of the landscape [MSDH19]. Figure 2 (left) contains an example of such a figure. Each solid black line corresponds to a change in gradient induced in the space by a neuron changing sign from positive to negative (or vice versa)—ignoring for now the remaining lines. The problem of neural network extraction corresponds to recovering the locations and angles of these neuron-induced hyperplanes: in general with input dimension $d_{0}$, the planes have dimension $d_{0}-1$.

The function that computes the first $j$ layers (up to and including $f_{j}$ but not including $\sigma$) of $f$ is denoted as $f_{1..j}$. In particular, $f=f_{1..k}$.

The hidden state at layer $j$ is the output of the function $f_{1..j}$, before applying the nonlinear transformation $\sigma$.

Layer $f_{j}$ is a linear transformation of the $(j-1)$st hidden state after $\sigma$.

$\mathcal{V}(\eta;x)$ denotes the input to neuron $\eta$ (before applying $\sigma$) when evaluated at $x$. $\mathcal{L}(\eta)$ denotes the layer of neuron $\eta$. The first layer starts at 1.

A neuron $\eta$ is at a critical point when $\mathcal{V}(\eta;x)=0$. We refer to this input $x$ as a witness to the fact that $\eta$ is at a critical point, denoted by $x\in\mathcal{W}(\eta)$. If $\mathcal{V}(\eta;x)>0$ then $\eta$ is active, and otherwise inactive.

In Figure 2 the locations of these critical points correspond exactly to the solid black lines drawn through the plane. Observe that because we restrict ourselves to ReLU neural networks, the function $f$ is piecewise linear and infinitely differentiable almost everywhere. The gradient $\nabla_{x}f(x)$ is well defined at all points $x$ except when there exists a neuron that is at its critical point.

Functionally, the attack as presented in this subsection has appeared previously in the literature [MSDH19, JCB+20]. By framing it differently, our attack will be extensible to deeper networks.

Assume we were given a witness $x^{*}\in\mathcal{W}(\eta_{j})$ that caused neuron $\eta_{j}$ to be at its critical point (i.e., its value is identically zero). Because we are using the ReLU activation function, this is the point at which that neuron is currently “inactive” (i.e., is not contributing to the output of the classifier) but would become “active” (i.e., contributing to the output) if it becomes slightly positive. Further assume that only this neuron $\eta_{j}$ is at its critical point, and that for all others neurons $\eta\neq\eta_{j}$ we have $|\mathcal{V}(\eta,x_{j})|>\delta$ for a constant $\delta>0$.

Consider the quantity $\lvert\alpha_{+}-\alpha_{-}\rvert$. Because $x^{*}$ induces a critical point of $\eta_{j}$, exactly one of $\{\alpha_{+},\alpha_{-}\}$ will have the neuron $\eta_{j}$ in its active regime and the other will have $\eta_{j}$ in its inactive regime. If no two columns of $A^{(1)}$ are collinear, then as long as $\epsilon<{\delta\over\sum_{i,j}\lvert A^{(1)}_{i,j}\rvert}$, we are guaranteed that all other neurons in the neural network will remain in the same state as before—either active or inactive. Therefore, if we compute the difference $\lvert\alpha^{i}_{+}-\alpha^{i}_{-}\rvert$, the gradient information flowing into and out of all other neurons will cancel and we will be left with just the gradient information flowing along the edge from the input coordinate $i$ to neuron $\eta_{j}$ to the output. Concretely, we can write the 1-deep neural network as

and so either $\alpha^{i}_{+}-\alpha^{i}_{-}=A^{(1)}_{j,i}\cdot A^{(2)}$ or $\alpha^{i}_{-}-\alpha^{i}_{+}=A^{(1)}_{j,i}\cdot A^{(2)}$. However, if we repeat the above procedure on a new basis vector $e_{k}$ then either $\alpha^{k}_{+}-\alpha^{k}_{-}=A^{(1)}_{j,k}\cdot A^{(2)}$ or $\alpha^{k}_{-}-\alpha^{k}_{+}=A^{(1)}_{j,k}\cdot A^{(2)}$ will hold. Crucially, whichever of the two relations that holds for along coordinate $i$ will be the same relation that holds on coordinate $k$. Therefore we can divide out $A^{(2)}$ to obtain the ratio of pairs of weights

This allows us to compute every row of $A^{(1)}$ up to a single scalar $c_{j}$. Further, we can compute $b^{(1)}_{j}=-\hat{A}^{(1)}_{j}\cdot x^{*}$ (again, up to a scaling factor) because we know that $x^{*}$ induces a critical point on neuron $\eta_{j}$ and so its value is zero.

Observe that the magnitude of $c_{j}$ is unimportant. We can always push a constant $c>0$ through to the weight matrix $A^{(2)}$ and have a functionally equivalent result. However, the sign of $c_{j}$ does matter.

Because $x_{i}$ is a witness to neuron $\eta_{i}$ being at its critical point, we will have that either $f(x_{+})=f(x_{i})$ or $f(x_{-})=f(x_{i})$. Exactly one of these equalities is true because $\sigma(h-e_{i})=\sigma(h)$, but $\sigma(h+e_{i})\neq\sigma(h)$ when $h_{i}=0$. Therefore if the second equality holds true, then we know that our extracted guess of the $i$th row has the correct sign. However, if the first equality holds true, then our extracted guess of the $i$th row has the incorrect sign, and so we invert it (along with the bias $b^{(1)}_{i}$). We repeat this procedure with a critical point for every neuron $\eta_{i}$ to completely recover the signs for the full first layer.

It only remains to show how to find witnesses $x^{*}\in\mathcal{W}(\eta)$ for each neuron $\eta$ on the first layer. We choose a random line in input space (the dashed line in Figure 2, left), and search along it for nonlinearities in the partial derivative. Any nonlinearity must have resulted from a ReLU changing signs, and locating the specific location where the ReLU changes signs will give us a critical point. We do this by binary search.

Once we have fully recovered the first layer weights, we can “peel off” the weight matrix $A^{(1)}$ and bias $b^{(1)}$ and we are left with extracting the final linear layer, which reduces to -deep extraction.

3 k𝑘k-Deep Contractive Neural Networks

Extending the above attack to deep neural networks has several complications that prior work was unable to resolve efficiently; we address them one at a time.

Because 1-deep networks have only one layer, all ReLUs occur on that layer. Therefore all critical points found during search will correspond to a neuron on that layer. For $k$-deep networks this is not true, and if we want to begin by extracting the first layer we will have to remove non-first layer critical points. (And, in general, to extract layer $j$, we will have to remove non-layer-$j$ critical points.)

In order to be able to directly read off the weights, we query the network on basis vectors $e_{i}$. Achieving this is not always possible for deep networks, and we must account for the fact that we may only be able to query on non-orthogonal directions.

Our row-sign procedure requires that we be able to invert $A^{(1)}$, which in general implies we need to develop a method to compute a preimage of $f_{1..j}$.

3.1 Extracting layer-1 weights with unknown critical point layers

Suppose we had a function $\mathcal{C}_{0}(f)=\{x_{i}\}_{i=1}^{M}$ that returns at least one critical point for every neuron in the first layer (implying $M\geq d_{1}$), but never returns critical points for any deeper layer. We claim that the exact differential attack from above still correctly recovers the first layer of a deep neural network.

We make the following observation. Let $x^{*}\not\in\bigcup_{\eta_{i}}\mathcal{W}(\eta_{i})$ be an input that is a witness to no critical point, i.e., $\lvert\mathcal{V}(\eta_{i};x^{*})\rvert>\epsilon>0$. Define $f_{\text{local}}$ as the function so that for a sufficiently small region we have that $f_{\text{local}}\equiv f$, that is,

Here, $I^{(j)}$ are 0-1 diagonal matrices with a on the diagonal when the neuron is inactive and $1$ on the diagonal when the neuron is active:

where $\eta_{n}$ is the $n$th neuron on the first layer. Importantly, observe that each $I^{(j)}$ is a constant as long as $x$ is sufficiently close to $x^{*}$. While $\beta$ is unknown, as long as we make only gradient queries $\partial f_{\text{local}}$, its value is unimportant. This observation so far follows from the definition of piecewise linearity.

Consider now some input that is a witness to exactly one critical point on neuron $\eta^{*}$. Formally, $x^{*}\in\mathcal{W}(\eta^{*})$, but $x^{*}\not\in\bigcup_{\eta_{j}\neq\eta^{*}}\mathcal{W}(\eta_{j};x^{*})$. Then

where again $I^{(j)}$ are 0-1 matrices, but except that now, $I^{(1)}$ (and only $I^{(1)}$) is a function of $x$ returning a 0-1 diagonal matrix that has one of two values, depending on the value of $\mathcal{V}(\eta^{*};x)>0$. Therefore we can no longer collapse the matrix product into one matrix $\Gamma$ but instead can only obtain

But this is exactly the case we have already solved for 1-deep neural network weight recovery: it is equivalent to the statement $f_{\text{local}}(x)=\Gamma\sigma(A^{(1)}x+b^{(1)})+\beta_{2}$, and so by dividing out $\Gamma$ exactly as before we can recover the ratios of $A^{(1)}_{i,j}$.

Assume we are given a set of inputs $S=\{x_{i}\}$ so that each $x_{i}$ is a witness to neuron $\eta_{x_{i}}$, with $\eta_{x_{i}}$ unknown. By the coupon collector’s argument (assuming uniformity), for $|S|\gg N\log N$, where $N$ is the total number of neurons, we will have at least two witnesses to every neuron $\eta$.

Without loss of generality let $x_{0},x_{1}\in\mathcal{W}(\eta)$ be witnesses to the same neuron $\eta$ on the first layer, i.e, that $\mathcal{V}(\eta;x_{0})=\mathcal{V}(\eta;x_{1})=0$. Then, performing the weight recovery procedure beginning from each of these witnesses (through finite differences) will yield the correct weight vector $A^{(1)}_{j}$ up to a scalar.

Typically elements of $S$ will not be witnesses to neurons on the first layer. Without loss of generality let $x_{2}$ and $x_{3}$ be witnesses to any neuron on a deeper layer. We claim that we will be able to detect these error cases: the outputs of the extraction algorithm will appear to be random and uncorrelated. Informally speaking, because we are running an attack designed to extract first-layer neurons on a neuron actually on a later layer, it is exceedingly unlikely that the attack would, by chance, give consistent results when run on $x_{2}$ and $x_{3}$ (or any arbitrary pair of neurons).

Formally, let $h_{2}=f_{1}(x_{2})$ and $h_{3}=f_{1}(x_{3})$. With high probability, $\text{sign}(h_{2})\neq\text{sign}(h_{3})$. Therefore, when executing the extraction procedure on $x_{2}$ we compute over the function $\Gamma_{1}I^{(1)}(x_{2})A^{(1)}x+\beta_{1}$, whereas extracting on $x_{3}$ computes over $\Gamma_{2}I^{(1)}(x_{3})A^{(1)}x+\beta_{2}$. Because $\Gamma_{1}\neq\Gamma_{2}$, this will give inconsistent results.

Therefore our first layer weight recovery procedure is as follows. For all inputs $x_{i}\in S$ run the weight recovery procedure to recover the unit-length normal vector to each critical hyperplane. We should expect to see a large number of vectors only once (because they were the result of running the extraction of a layer $2$ or greater neuron), and a small number of vectors that appear duplicated (because they were the result of successful extraction on the first layer). Given the first layer, we can reduce the neural network from a $k$-deep neural network to a $(k-1)$-deep neural network and repeat the attack. We must resolve two difficulties, however, discussed in the following two subsections.

3.2 Extracting hidden layer weights with unknown critical points

When extracting the first layer weight matrix, we were able to compute $\partial^{2}f\over\partial e_{1}\partial e_{j}$ for each input basis vectors $e_{i}$, allowing us to “read off” the ratios of the weights on the first layer directly from the partial derivatives. However, for deeper layers, it is nontrivial to exactly control the hidden layers and change just one coordinate in order to perform finite differences.For the expansive networks we will discuss in Section 4.4 it is actually impossible; therefore this section introduces the most general method. Let $j$ denote the current layer we are extracting. Begin by sampling $d_{j}+1$ directions $\delta_{i}\sim\mathcal{N}(0,\epsilon I_{d_{0}})\in\mathcal{X}$ and let

From here we can construct a system of equations: let $h_{i}=\sigma(f_{1..j-1}(x+\delta_{i}))$ and solve for the vector $w$ such that $h_{i}\cdot w=y_{i}$.

As before, we run the weight recovery procedure assuming that each witness corresponds to a critical point on the correct layer. Witnesses that correspond to neurons on incorrect layers will give uncorrelated errors that can be discarded.

The above algorithm overlooks one important problem. For a given critical point $x^{*}$, the hidden vector obtained from $f_{1..j}(x^{*})$ is likely to have several (on average, half) neurons that are negative, and therefore $\sigma(f_{1..j}(x^{*}))$ and any $\sigma(f_{1..j}(x^{*}+\delta_{i}))$ will have neurons that are identically zero. This makes it impossible to recover the complete weight vector from just one application of least squares—it is only possible to compute the weights for those entries that are non-zero. One solution would be to search for a witness $x^{*}$ such that component-wise $f_{1..j}(x^{*})\geq 0$; however doing this is is not possible in general, and so we do not consider this option further.

Instead, we combine together multiple attempts at extracting the weights through a unification procedure. If $x_{1}$ and $x_{2}$ are witnesses to critical points for the same neuron, and the partial vector $f_{1..j}(x_{1})$ has entries $t_{1}\subset\{1,\dots,d_{j}\}$ and the partial vector $f_{1..j}(x_{2})$ has entries $t_{2}\subset\{1,\dots,d_{j}\}$ defined, then it is possible to recover the ratios for all entries $t_{1}\cup t_{2}$ by unifying together the two partial solutions as long as $t_{1}\cap t_{2}$ is non-empty as follows.

Let $r_{i}$ denote the extracted weight vector on witness $x_{1}$ with entries at locations $t_{1}\subset\{1,\dots,d_{j}\}$ (respectively, $r_{2}$ at $x_{2}$ with locations at $t_{2}$). Because the two vectors correspond to the solution for the same row of the weight matrix $A^{(j)}_{i}$, the vectors $r_{1}$ and $r_{2}$ must be consistent on $t_{1}\cap t_{2}$. Therefore, we will have that $r_{1}[t_{1}\cap t_{2}]=c\cdot r_{2}[t_{1}\cap t_{2}]$ for a scalar $c\neq 0$. As long as $t_{1}\cap t_{2}\neq\emptyset$ we can compute the appropriate constant $c$ and then recover the weight vector $r_{1,2}$ with entries at positions $t_{1}\cup t_{2}$.

Observe that this procedure also allows us to check whether or not $x_{1}$ and $x_{2}$ are witnesses to the same neuron $n$ reaching its critical point. If $|t_{1}\cap t_{2}|\geq\gamma$, then as long as there do not exist two rows of $A^{(j)}$ that have $\gamma+1$ entries that are scalar multiples of each other, there will be a unique solution that merges the two partial solutions together. If the unification procedure above fails—because there does not exist a single scalar $c$ so that $c\cdot r_{1}[t_{1}\cap t_{2}]=r_{2}[t_{1}\cap t_{2}]$—then $x_{1}$ and $x_{2}$ are not witnesses to the same neuron being at a critical point.

3.3 Recovering row signs in deep networks

In the case of $1$-deep networks, it suffices for $d_{1}\leq d_{0}$ and $A^{(1)}$ to be onto as described. In general it is necessary that $d_{k}\leq d_{k-1}\leq\dots\leq d_{1}\leq d_{0}$ but it is not sufficient, even if every layer $A^{(i)}$ were an onto map. Because there is a ReLU activation after every hidden layer, it is not possible to send negative values into the second layer $f_{j}$ when computing the preimage.

It is simple to test if a network is sufficiently contractive without any queries: try the above method to find a preimage $x$; if this fails, abort and attempt the following (more expensive) attack procedure. Otherwise it is contractive.

4 k𝑘k-Deep Expansive Neural Networks

While most small neural networks are contractive, in practice almost all interesting neural networks are expansive: the number of neurons on some intermediate layer is larger than the number of inputs to that layer. Almost all of the prior methods still apply in this setting, with one exception: the column sign recovery procedure. Thus, we are required to develop a new strategy.

Observe that sign information is not lost for the final layer: because there is no ReLU activation and we can directly solve for the weights with least squares, we do not lose sign information.

Suppose we had extracted completely the function $\hat{f}_{1..k-1}$ (the third to last layer), and further had extracted the weights $\hat{A}^{(k)}$ and biases $\hat{b}^{(k)}$ up to sign of the rows. There are three unknown quantities remaining: a sign vector $s\in\{-1,1\}^{d_{k}}$, $\hat{A}^{(k+1)}$ and $\hat{b}^{(k+1)}$. Suppose we were given $S\subset\mathcal{X}$ so that $\lvert S\rvert>d_{k}$. Then it would be possible to solve for all three unknown simultaneously through brute force.

Let $h_{i}=\sigma(f_{1..k-1}(x_{i}))$. Enumerate all $2^{d_{k}}$ assignments of $s$ and compute $g_{i}=\sigma((s\odot\hat{A}^{(k)})h_{i}+(s\odot\hat{b}^{(k)})$. We know that if we guessed the sign vector $s$ correctly, then there would exist a solution to the system of equations $v\cdot g_{i}+b=f(x_{i})$. This is the zero-deep extraction problem and solving it efficiently requires just a single call to least squares. This allows us to—through brute forcing the sign bits—completely recover both the signs of the second-to-last layer as well as the values (and signs) of the final layer.

Unfortunately, this procedure does not scale to recover the signs of layer $k-1$ and earlier. It relies on the existence of an efficient testing procedure (namely, least squares) to solve the final layer. If we attempted this brute-force strategy at layer $k-3$ in order to test if our sign assignment was correct, we would need to run the complete layer $k-2$ extraction procedure, thus incurring an exponential number of queries to the oracle.

However, we can use this idea in order to still recover signs even at earlier layers in the network with only a linear number of queries (but still exponential work in the width of the hidden layers).

Assume that we are given a collection of examples $\{x_{i}\}\subset\mathcal{W}(\eta)$ for some neuron $\eta$ that is on the layer after we extracted so far: $\mathcal{L}(\eta)=j+1$. Then we would know that there should exist a single unknown vector $v$ and bias $b$ such that $f_{j}(x_{i})\cdot v+b=0$ for all $x_{i}$.

This gives us an efficient procedure to test whether or not a given sign assignment on layer $j$ is correct. As before, we enumerate all possible sign assignments and then check if we can recover such a vector $v$. If so, the assignment is correct; if not, it is wrong. It only remains left to show how to implement this procedure to obtain such a collection of inputs $\{x_{i}\}$.

4.1 The polytope boundary projection algorithm

The layer $j$ polytope containing $x$ is the set of points $\{x+\delta\}$ so that $\text{sign}(\mathcal{V}(\eta;x))=\text{sign}(\mathcal{V}(\eta;x+\delta))$ for all $\mathcal{L}(\eta)\leq j$.

Observe that the layer $j$ polytope around $x$ is an open, convex set, as long as $x$ is not a witness to a critical point. In Figure 3, each enclosed region is a layer-$k$ polytope and the triangle formed by $\eta_{0},\eta_{1}$, and $\eta_{2}$ is a layer-$(k-1)$ polytope.

Given an input $x$ and direction $\Delta$, we can compute the distance $\alpha$ so that the value $x^{\prime}=x+\alpha\Delta$ is at the boundary of the polytope defined by layers $1$ to $k$. That is, starting from $x$ traveling along direction $\Delta$ we stop the first time a neuron on layer $j$ or earlier reaches a critical point. Formally, we define

We only ever compute $\text{Proj}_{1..j}$ when we have extracted the neural network up to layer $j$. Thus we perform the computation with respect to the extracted function $\hat{f}$ and neuron-value function $\hat{\mathcal{V}}$, and so computing this function requires no queries to the oracle. In practice we solve for $\alpha$ via binary search.

4.2 Identifying a single next-layer witness

Given the correctly extracted network $\hat{f}_{1..j-1}$ and the weights (up to sign) of layer $j-1$, our sign extraction procedure requires some witness to a critical point on layer $j$. We begin by performing our standard binary search sweep to find a collection $S\subset\mathcal{X}$, each of which is a witness to some neuron on an unknown layer. It is simple to filter out critical points on layers $j-1$ or earlier by checking if any of $\hat{\mathcal{V}}(\eta;x)=0$ for $\mathcal{L}(\eta)\leq j-1$. Even though we have not solved for the sign of layer $j$, it is still possible to compute whether or not they are at a critical point because critical points of $\hat{A}^{(j)}$ are critical points of $-\hat{A}^{(j)}$. This removes any witnesses to critical points on layer $j$ or lower.

Now we must filter out any critical points on layers strictly later than $j$. Let $x^{*}\in\mathcal{W}(\eta^{*})$ denote a potential witness that is on layer $j$ or later (having already filtered out critical points on layers $j-1$ or earlier). Through finite differences, estimate $g=\pm\nabla_{x}f(x)$ evaluated at $x=x^{*}$. Choose any random vector $r$ perpendicular to $g$, and therefore parallel to the critical hyperplane. Let $\alpha=\text{Proj}_{1..j}(x^{*},r)$. If it turns out that $x^{*}$ is a witness to a critical point on layer $j$ then for all $\epsilon<\alpha$ we must have that $x^{*}+\epsilon r\in\mathcal{W}(\eta^{*})$. Importantly, we also have the converse: with high probability for $\delta>\alpha$ we have that $x^{*}+\delta r\not\in\mathcal{W}(\eta^{*})$. However, observe that if $x^{*}$ is not a witness to a neuron on layer $j$ then one of these two conditions will be false. We have already ruled out witnesses on earlier neuron, so if $x^{*}$ is a witness to a later neuron on layer $j^{\prime}>j$ then it is unlikely that the layer-$j^{\prime}$ polytope is the same shape as the layer-$j$ polytope, and therefore we will discover this fact. In the case that the two polytopes are actually identical, we can mount the following attack and if it fails we know that our initial input was on the wrong layer.

4.3 Recovering multiple witnesses for the same neuron

The above procedure yields a single witness $x^{*}\in\mathcal{W}(\eta^{*})$ so that $\mathcal{L}(\eta^{*})=j+1$. We expand this to a collection of witnesses $W$ where all $x\in W$ have $x\in\mathcal{W}(\eta^{*})$, requiring the set to be diverse:

A collection of inputs $S$ is fully diverse at layer $j$ if for all $\eta$ with $\mathcal{L}(\eta)=j$ and for $s\in\{-1,1\}$ there exists $x\in S$ such that $s\cdot\mathcal{V}(\eta;x)\geq 0$.

Informally, being diverse at layer $j$ means that if we consider the projection onto the space of layer $j$ (by computing $f_{1..j}(x)$ for $x\in S$), for every neuron $\eta$ there will be at least one input $x_{+}\in S$ that the neuron is positive, and at least one input $x_{-}\in S$ so that the neuron is negative.

Our procedure is as follows. Let $n$ be normal to the hyperplane $x^{*}$ is on. Choose some $r$ with $r\cdot n=0$ and let $\alpha=\text{Proj}_{1..j}(x^{*},r)$ to define $x^{\prime}=x^{*}+\alpha r$ as a point on the layer-$j$ polytope boundary. In particular, this implies that we still have that $x^{\prime}\in\mathcal{W}(\eta^{*})$ (because $r$ is perpendicular to $n$) but also $x^{\prime}\in\mathcal{W}(\eta_{u})$ for some neuron $\mathcal{L}(\eta_{u})<j$ (by construction of $\alpha$). Call this input $x^{\prime}$ the double-critical point (because it is a witness to two critical points simultaneously).

From this point $x^{\prime}$, we would like to obtain a new point $y$ so that we still have $y\in\mathcal{W}(\eta^{*})$, but that also $y$ is on the other side of the neuron $\eta_{u}$, i.e., $\text{sign}(\mathcal{V}(\eta_{u};x^{*}))\neq\text{sign}(\mathcal{V}(\eta_{u};y))$. Figure 3 (right) gives a diagram of this process. In order to follow $x^{*}$ along its path, we first need to find it a critical point on the new hyperplane, having just been bent by the neuron $\eta_{u}$. We achieve this by performing a critical-point search starting $\epsilon$-far away from, and parallel to, the hyperplane from neuron $\eta_{u}$ (the dashed line in Figure 3). This returns a point $y$ from where we can continue the hyperplane following procedure.

The geometric view hurts us here: because the diagram is a two-dimensional projection, it appears that from the critical point $y$ there are only two directions we can travel in: away from $x^{\prime}$ or towards $x^{\prime}$. Traveling away is preferable—traveling towards $x^{\prime}$ will not help us construct a fully diverse set of inputs.

However, a $d_{0}$-dimensional input space has in general $(d_{0}-1)$ dimensions that remain on the neuron $\eta^{*}$. We defer to Section 5.5.1 an efficient method for selecting the continuation direction. For now, observe that choosing a random direction will eventually succeed at constructing a fully-diverse set, but is extremely inefficient: there exist better strategies than choosing the next direction.

4.4 Brute force recovery

Given this collection $S$, we can now—through brute force work—recover the correct sign assignment as follows. As described above, compute a fully diverse set of inputs $\{x_{i}\}$ and define $h_{i}=f_{1..j}(x_{i})$. Then, for all possible $2^{d_{j}}$ assignments of signs $s\in\{-1,1\}^{d_{j}}$, compute the guessed weight matrix $\hat{A}^{(j)}_{s}=s\odot\hat{A}^{(j)}$.

Instantiating the Differential Attack in Practice

The above idealized attack would efficiently extract neural network models but suffers from two problems. First, many of the algorithms are not numerically stable and introduce small errors in the extracted weights. Because errors in layer $i$ compound and cause further errors at layers $j>i$, it is necessary to keep errors to a minimum. Second, the attack requires more chosen-inputs than is necessary; we develop new algorithms that require fewer queries or re-use previously-queried samples.

Each sub-section that follows is independent from the surrounding sub-sections and modifies algorithms introduced in Section 4. For brevity, we assume complete knowledge of the original algorithm and share the same notation. Readers may find it helpful to review the original algorithm before proceeding to each subsection.

1 Improving Precision of Extracted Layers

To begin, select a neuron $\eta$ with $\mathcal{L}(\eta)=j$. By querying the already-extracted model $\hat{f}_{1..j}$, analytically compute witnesses $\{x_{i}\}_{i=1}^{d_{j}}$ so that each $x_{i}\in\hat{\mathcal{W}}(\eta)$. This requires no queries to the model as we have already extracted this partial model.

If the $\hat{A}^{(j)}$ and $\hat{b}^{(j)}$ were exactly correct then $\mathcal{W}(\eta;\cdot)\equiv\hat{\mathcal{W}}(\eta;\cdot)$ and so each computed critical point $x_{i}$ would be exactly a critical point of the true model $f$ and so $\mathcal{V}(\eta;x_{i})\equiv 0$. However, if there is any imprecision in the computation, then in general we will have that $0<\lvert\mathcal{V}(\eta;x_{i})\rvert<\epsilon$ for some small $\epsilon>0$.

Fortunately, given this $x_{i}$ it is easy to compute $x^{\prime}_{i}$ so that $\mathcal{V}(\eta;x^{\prime}_{i})=0$. To do this, we sample a random $\Delta\in\mathcal{R}^{d_{0}}$ and apply our binary search procedure on the range $[x_{i}+\Delta,x_{i}-\Delta]$. Here we should select $\Delta$ so that $\lVert\Delta\rVert$ is sufficiently small that the only critical points it crosses is the one induced by neuron $\eta$, but sufficiently large that it does reliably find the true critical point of $\eta$.

Repeating this procedure for each witness $x_{i}$ gives a set of witnesses $\{x^{\prime}_{i}\}_{i=1}^{d_{j}}$ to the same neuron $\eta$. We compute $h_{i}=\hat{f}_{1..j-1}(x^{\prime}_{i})$ as the hidden vector that layer $j$ will receive as input. By assumption $h_{i}$ is precise already and so $\hat{f}_{1..j-1}\approx f_{1..j-1}$. Because $x^{\prime}_{i}$ is a witness to neuron $\eta$ having value zero, we know that that $A^{(j)}_{n}\cdot h_{i}=0$ where $n$ corresponds to the row of neuron $\eta$ in $A^{(j)}$.

Ideally we would solve this resulting system with least squares. However, in practice, occasionally the conversion from $x\to x^{\prime}$ fails because $x^{\prime}$ is no longer a witness to the same neuron $\eta^{\prime}$. This happens when there is some other neuron (i.e., $\eta^{\prime}$) that is closer to $x$ than the true neuron $\eta$. Because least squares is not robust to outliers this procedure can fail to improve the solution.

We take two steps to ensure this does not happen. First, observe that if $\Delta$ is smaller, the likelihood of capturing incorrect neurons $\eta^{\prime}$ decreases faster than the likelihood of capturing the correct neuron $\eta$. Thus, we set $\Delta$ to be small enough that roughly half of the attempts at finding a witness $x^{\prime}$ fails. Second, we apply a (more) robust method of determining the vector that satisfies the solution of equations [JOB+18]. However, even these two techniques taken together occasionally fail to find valid solutions to improve the quality. When this happens, we reject this proposed improvement and keep the original value.

2 Efficient Finite Differences

Most of the methods in this paper are built on computing second partial derivatives of the neural network $f$, and therefore developing a robust method for estimating the gradient is necessary. Throughout Section 4 we compute the partial derivative of $f$ along direction $\alpha$ evaluated at $x$ with step size $\varepsilon$ as

To compute the second partial derivative earlier, we computed $\alpha^{i}_{+}$ and $\alpha^{i}_{-}$ by first taking a step towards $x^{*}+\epsilon_{0}e_{1}$ for a different step size $\epsilon_{0}$ and then computed the first partial derivative at this location. However, with floating point imprecision it is not desirable to have two step sizes ($\epsilon_{0}$ controlling the distance away from $x^{*}$ to step, and $\epsilon$ controlling the step size when computing the partial derivative). Worse, we must have that $\epsilon\ll\epsilon_{0}$ because if ${\partial f\over\partial e_{1}}\epsilon_{0}>{\partial f\over\partial e_{i}}\epsilon$ then when computing the partial derivative along $e_{i}$ we may cross the hyperplane and estimate the first partial derivative incorrectly. Therefore, instead we compute

where we both step along $e_{i}$ and also take the partial derivative along the same $e_{i}$ (and similarly for -$e_{i}$). This removes the requirement for an additional hyperparameter and allows the step size $\epsilon$ to be orders of magnitude larger, but introduces a new error: we now lose the relative signs of the entries in the row when performing extraction and can only recover $\left\rvert A^{(1)}_{i,j}/A^{(1)}_{i,k}\right\rvert$.

We next recover the value $\text{sign}(A^{(1)}_{i,j})\cdot\text{sign}(A^{(1)}_{i,k})$. Fortunately, the same differencing process allows us to learn this information, using the following observation: if $A^{(1)}_{i,j}$ and $A^{(1)}_{i,k}$ have the same sign, then moving in the $e_{j}+e_{k}$ direction will cause their contributions to add. If they have different signs, their contributions will cancel each other. That is, if

We can repeat this process to test whether each $A^{(1)}_{i,j}$ has the same sign as (for example) $A^{(1)}_{i,1}$. However, we still do not know whether any single $A^{(1)}_{i,j}$ is positive or negative—we still must recover the row signs as done previously.

3 Finding Witnesses to Critical Points

Throughout the paper we require the ability to find witnesses to critical points. Section 4.2 uses simple binary search to achieve this which is (a) imprecise in practice, and (b) query inefficient. We improve on the witness-finding search procedure developed by [JCB+20]. Again we interpolate between two examples $u,v$ and let $x_{\alpha}=(1-\alpha)u+\alpha v$. Previously, we repeatedly performed binary search as long as the partial derivatives were not equal $\partial f(x_{\alpha})\neq\partial f(x_{\beta})$, requiring $p$ queries to obtain $p$ bits of precision of the value $x^{*}$. However, observe that if $x_{\alpha}$ and $x_{\beta}$ differ in the sign of exactly one neuron $i$, then we can directly compute the location $x^{*}$ at which $\mathcal{V}(\eta_{i};x^{*})=0$ but so that for all other $\eta_{j}$ we have

This approach is illustrated in Figure 4 and relies on the fact that $f$ is a piecewise linear function with two components. By measuring, $f(x_{\alpha})$ and $\partial f(x_{\alpha})$ (resp., $f(x_{\beta})$ and $\partial f(x_{\beta})$), we find the slope and intercept of both the left and right lines in Figure 4 (left). This allows us to solve for their expected intersection $(x^{*},\hat{f}(x^{*}))$. Typically, if there are more than two linear segments, as in the middle of the figure, we will find that the true function value $f(x^{*})$ will not agree with the expected function value $\hat{f}(x^{*})$ we obtained by computing the intersection; we can then perform binary search again and repeat the procedure.

However, we lose some soundness from this procedure. As we see in Figure 4 (right), situations may arise where many ReLU units change sign between $x_{\alpha}$ and $x_{\beta}$, but $\hat{f}(x^{*})=f(x^{*})$. In this case, we would erroneously return $x^{*}$ as a critical point, and miss all of the other critical points in the range. Fortunately, this error case is pathological and does not occur in practice.

Suppose that we had already extracted the first $j$ layers of the neural network and would like to perform the above critical-point finding algorithm to identify all critical points between $x_{\alpha}$ and $x_{\beta}$. Notice that we do not need to collect any more critical points from the first $j$ layers, but running binary search will recover them nonetheless. To bypass this, we can analytically compute $S$ as the set of all witnesses to critical points on the extracted neural network $\hat{f}_{1..j}$ between $x_{\alpha}$ and $x_{\beta}$. As long as the extracted network $\hat{f}$ is correct so far, we are guaranteed that all points in $S$ are also witnesses to critical points of the true $f$.

Instead of querying on the range $(x_{\alpha},x_{\beta})$ we perform the $\lvert S\rvert+1$ different searches. Order the elements of $S$ as $\{s_{i}\}_{i=1}^{|S|}$ so that $s_{i}<s_{j}\implies\lvert x_{\alpha}-s_{i}\rvert<\lvert x_{\alpha}-s_{j}\rvert$. Abusing notation, let $s_{1}=x_{\alpha}$ and $s_{|S|}=x_{\beta}$. Then, perform binary search on each disjoint range $(S_{i},S_{i+1})$ for $i=1$ to $|S|-1$ and return the union.

4 Unification of Witnesses with Noisy Gradients

Recall that to extract $\hat{A}^{(l)}$ we extract candidates candidates $\{r_{i}\}$ and search for pairs $r_{i},r_{j}$ that agree on multiple coordinates. This allows us to merge $r_{i}$ and $r_{j}$ to recover (eventually) full rows of $\hat{A}^{(l)}$. With floating point error, the unification algorithm in Section 4.3.2 fails for several reasons.

Our core algorithm computes the normal to a hyperplane, returning pairwise ratios $\hat{A}^{(1)}_{i,j}/\hat{A}^{(1)}_{i,k}$; throughout Section 4 we set $\hat{A}^{(1)}_{i,1}=1$ without loss of generality.

Unfortunately in practice there is loss of generality, due to the disparate impact of numerical instability. Consider the case where $A^{(l)}_{i,1}<10^{-\alpha}$ for $\alpha\gg 0$, but $A^{(l)}_{i,k}\geq 1$ for all other $k$. Then there will be substantially more (relative) floating point imprecision in the weight $A^{(l)}_{i,1}$ than in the other weights. Before normalizing there is no cause for concern since the absolute error is no larger than for any other. However, the described algorithm now normalizes every other coordinate $A^{(l)}_{i,k}$ by dividing it by $A^{(l)}_{i,1}$—polluting the precision of these values.

Therefore we adjust our solution. At layer $l$, we are given a collection of vectors $R=\{r_{i}\}_{i=1}^{n}$ so that each $r_{i}$ corresponds to the extraction of some (unknown) neuron $\eta_{i}$. First, we need an algorithm to cluster the items into sets $\{S_{j}\}_{j=1}^{d_{l}}$ so that $S_{j}\subset R$ and so that every vector in $S_{j}$ corresponds to one neuron on layer $l$. We then need to unify each set $S_{j}$ to obtain the final row of $\hat{A}^{(l)}_{j}$.

We construct the three dimensional $M_{i,a,b}=r^{(i)}_{a}/r^{(i)}_{b}$. Given $M$, the a good guess for the scalar $c_{ab}$ so that $r^{(i)}_{a}=r^{(i)}_{b}\cdot C_{ab}$ along as many coordinates $i$ as possible is the assignment $C_{ab}=\text{median}_{i}\,M_{i,a,b}$, where the estimated error is $e_{ab}=\text{stdev}_{i}\,M_{i,a,b}$.

If all $r_{a}$ were complete and had no imprecision then $C_{ab}$ would have no error and so $C_{ab}=C_{ax}\cdot C_{xb}$. However because it does have error, we can iteratively improve the guessed $C$ matrix by observing that if the error $e_{ax}+e_{xb}<e_{ab}$ then the guessed assignment $C_{ax}\cdot C_{xb}$ is a better guess than $C_{ab}$. Thus we replace $C_{ab}\leftarrow C_{ax}\cdot C_{xb}$ and update $e_{ab}\leftarrow e_{ax}+e_{xb}$. We iterate this process until there is no further improvement. Then, finally, we choose the optimal dimension $a=\mathop{\text{arg min}}_{a}\sum_{b}e_{ab}$ and return the vector $C_{a}$. Observe that this procedure closely follows constructing the union of two partial entries $r_{i}$ and $r_{j}$ except that we perform it along the best axis possible for each coordinate.

5 Following Neuron Critical Points

Section 4.4.3 developed techniques to construct a set of witnesses to the same neuron being at its critical point. We now numerically-stabilize this procedure.

As before we begin with an input $x^{*}\in\mathcal{W}(\eta^{*})$ and compute the normal vector $n$ to the critical plane at $x^{*}$, and then choose $r$ satisfying $r\cdot n=0$. The computation of $n$ will necessarily have some floating point error, so $r$ will too.

This means when we compute $\alpha=\text{Proj}_{1..j}(x^{*},r)$ and let $x^{\prime}=x^{*}+r\alpha$ the resulting $x^{\prime}$ will be almost exactly a witness to some neuron $\eta_{u}$ with $\mathcal{L}(\eta_{u})<j$, (because this computation was performed analytically on a precisely extracted model), but $x^{\prime}$ has likely drifted off of the original critical plane induced by $\eta^{*}$.

To address this, after computing $\alpha$ we initially take a smaller step and let $x_{1}=x^{*}+r\sqrt{\alpha}$. We then refine the location of this point to a point $x_{2}$ by performing binary search on the region $x_{1}-\epsilon n$ to $x_{1}+\epsilon n$ for a small step $\epsilon$. If there was no error in computing $n$ then $x_{1}=x_{2}$ because both are already witnesses to $\eta^{*}$. If not, any error has been corrected. Given $x^{*}$ and $x_{2}$ we now can now compute $\alpha_{2}=\text{Proj}_{1..j}(x^{*},x_{2}-x^{*})$ and let $\bar{x}^{\prime}=x^{*}+(x_{2}-x^{*})\alpha_{2}$ which will actually be a witness to both neurons simultaneously.

Next we give a stable method to compute $y$ that is a witness to $\eta^{*}$ and on the other side of $\eta_{u}$. The previous procedure required a search parallel to $\eta_{u}$ and infinitesimally displaced, but this is not numerically stable without accurately yet knowing the normal to the hyperplane given by $\eta_{u}$.