Estimating the Jacobian matrix of an unknown multivariate function from sample values by means of a neural network

We propose to use neural networks to approximate the Jacobian matrix $J$ of a multivariate function $F$, where in the training of the neural network we use only a finite sample of input-output pairs $(x,F(x))$. Crucially, no additional information about $F$ or $J$ enters into the training and thus, in essence, the method proposed here differentiates a multivariate function $F$ solely based on a cloud of sample points.

The algorithm is based on two main ideas:

• •

  a loss function that utilizes a linear approximation of the sampled function $F$ in terms of the sought-after Jacobian matrix,

• •

  a nearest neighbor search in preparation of the training data from samples.

The neural network is not differentiated during or at the conclusion of the training.

In general, to estimate $J$ from sample points of $F$ is a difficult mathematical problem. Since the Jacobian has numerous applications in mathematics, science and engineering, the ability to estimate it by any means, for instance by a neural network, is valuable in its own right.

Of particular importance here is the fact that the Jacobian can be invoked to detect relations, or lack thereof, between the input variables and output values of the unknown function $F$. While neural networks are a powerful tool for nonlinear regression, the resulting interpolating function is notoriously a black-box, revealing little structural information about the sampled function $F$. By taking advantage of the estimated Jacobian matrix of $F$, we can start peering inside $F$. We therefore expect that the ideas presented here will contribute to new approaches for the training of neural networks designed to reveal structural information about sampled data.

A natural first approach to computing the Jacobian of a function $F$ by a neural network is to train a network to interpolate $F$ as usual and then differentiate the network itself (which, in the end, is a function).

One drawback of this approach is that it is possible to differentiate a neural network only if all its activation functions are differentiable, a situation that complicates the usage of some popular activation functions, such as ReLu [7]. (See [2] for a way of handling activation functions that are differentiable almost everywhere.)

A more serious obstacle is the well-known fact that even if a given function is a good approximation of another function—say in the sense of the uniform norm—their derivatives can be drastically different [14]. A simple but illustrative example is given in Figure 1. The sequence of sine waves converges uniformly on $\mathbb{R}$ to the constant zero function since the amplitude of the waves decreases to $0$. However, as the frequencies of the waves grow to infinity, the derivatives of the sine waves grow arbitrarily large and do not converge at all (not even pointwise) over $\mathbb{R}$.

In the context of neural networks, if a neural network is overfitted, then, in a manner similar to Runge’s phenomenon [16], the derivative of the overfitted neural network will likely have little resemblance to the derivative of the sampled function $F$.

The difficulty with differentiating a neural network trained to interpolate a function is well demonstrated in [9], where, among other results, the authors train a neural network $A$ on a sampled function $F$, differentiate the resulting neural network as a function, thus obtaining its Jacobian matrix $J_{A}$, and then compare the interpolated values of $F$ produced by $A$ with linear approximations of $F$ based on the Jacobian matrix $J_{A}$. They find that “the $J$ estimation task fails” and conclude that neural networks “may be not suitable to handle derivative signal analysis.”

In the highly influential paper [11], the authors proposed a method for training a neural network to solve (higher order) differential equations. The main idea of [11] is to differentiate the neural network during training and use this (higher order) derivative in the loss function. The training points are typically selected on a regular grid and hence it must be possible to evaluate all relevant functions at such points. The method of [11] is robust and it can be adopted to multivariate functions. We mention [13] as one of the many papers ultimately based on [11].

In the simplest case covered by [11], an approximate solution $y$ to the differential equation $y^{\prime}(x)=F(x)$ is obtained by comparing the neural network derivative $A^{\prime}$ to $F$ during training. Once $A$ is fully trained, $A^{\prime}$ should be a good approximation of $F$, that is, $A$ itself should be a good approximation of an antiderivate of $F$.

In order to obtain an approximation for the derivative of $F$, one could now differentiate the trained network $A$ twice (running into the issues described in the previous subsection), or start over with the differential equation $y^{\prime}(x)=F^{\prime\prime}(x)$ (which presumably requires $F^{\prime}$ to be already known), or consider a variation of [11] in which the neural network $A$ is integrated rather than differentiated during training. Little seems to be known about integration of neural networks.

One might work with data sampled from a function that satisfies a known differential equation, say for physical reasons. In the simplest case where the equation is first order and linear, the Jacobian matrix of the function can thus be computed directly from the differential equation. This essentially avoids the problem of estimating the Jacobian matrix by neural network altogether and relies on standard mathematical methods.

Similarly, the well-known automatic differentiation tool Autograd [12] (that is often used in connection with neural networks) conveniently differentiates a function whose code is written in NumPy [8], but it requires the function to be given explicitly (in NumPy code), not just by sample values.

Gradient estimation is a vast topic, cf. [6]. In the context of neural networks, gradient estimation typically refers to an estimation of the gradient of the explicitly given loss function. Gradient estimation is critical for the stochastic gradient descent algorithm, which is in turn key in the learning algorithm of neural networks [15]. The backpropagation algorithm [3, 5, 10] is, in essence, an algorithm that efficiently computes the gradient of a loss function using a graph-directed implementation of the chain rule. However, backpropagation relies on symbolic differentiation or on a tool similar to Autograd, as well as on a specific form of the neural network.

Another meaning of the phrase “gradient estimation” arises in situations when differentiating a given function is slow or impossible and its Jacobian matrix is therefore merely approximated. For instance, in [17] the authors give an efficient approximation of the Jacobian matrix by fast Fourier transform in the context of MRI. Such methods typically rely on a prescribed model for the function whose gradient is being estimated.

In contrast with all these results, we do not differentiate a neural network that has been trained to interpolate $F$, nor do we differentiate the neural network during training. Rather, we directly train a neural network to estimate the Jacobian matrix of an unknown function $F$ from given sample values of $F$, regardless of the context. We prove rigorously that the Jacobian estimator approaches the Jacobian of $F$ in norm, under reasonable assumptions on $F$, the sample set and the performance of the neural network used.

Suppose that $c$, $d$ are positive integers, $U$ is an open, bounded subset of $\mathbb{R}^{d}$, and $F:U\to\mathbb{R}^{c}$ is a Fréchet differentiable function on $U$ [14, Ch. 9]. (Note that $d$ stands for the dimension of the domain $\mathbb{R}^{d}$ and $c$ for the dimension of the codomain $\mathbb{R}^{c}$.)

Writing $x=(x_{1},\ldots,x_{d})$ for $x\in\mathbb{R}^{d}$ and $F=(F_{1},\ldots,F_{c})$ with $F_{j}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ for each $j\in\{1,\ldots,c\}$, the Jacobian matrix of $F$ is the matrix

In words, $J_{F}$ is the matrix whose rows are the gradients of $F_{1},\dots,F_{c}$.

Since the Jacobian $J_{F}$ is the matrix of the Fréchet derivative of $F$ with respect to the canonical basis, it satisfies by definition the following general linear approximation property (which is of course the very idea of the derivative):

where $\left\|{h}\right\|_{\mathbb{R}^{d}}$ is the usual Euclidean norm

Equation (III.1) will play a crucial role in the loss function of our neural network.

In a regression problem with noise, we are given a finite set $S=\left\{(x,F(x)+\epsilon_{x}):x\in X\right\}$ of pairs of input-output values sampled from an unknown function $F:U\subseteq\mathbb{R}^{d}\rightarrow\mathbb{R}^{c}$, where $X\subseteq U$ and $(\epsilon_{x})_{x\in X}$ is a family of independent random variables, all with mean $0$, representing noise. The goal is to find an estimate for $F$.

If we make the very strong assumption that $F$ is linear, i.e., that there exists a $c\times d$ matrix $A\in M_{c\times d}(\mathbb{R})$ such that $F(x)=Ax$ for all $x\in U$, then we may apply the standard linear, least square statistical method to derive an estimate for $A$, by minimizing the error

In this case, we note that the Jacobian matrix of $F$ is again $A$ and, of course, $A$ contains a lot of useful information about $F$.

Without the assumption of linearity of $F$, the problem becomes significantly more complicated and computationally intensive. A general method is to replace the algebra $M_{c\times d}(\mathbb{R})$ of matrices by a set of nonlinear functions which is parametrized by some points in the parameter space $\mathbb{R}^{P}$, where $P$ is typically quite large. The main example of interest here is a class of functions called artificial neural networks, which, in their simplest form, can be described as a finite chain of alternating compositions of linear functions and functions called activation functions.

Formally, an artificial neural network $A$ with $\ell$ layers, with $n_{0}$ inputs, and with $n_{j}$ neurons and activation function $f_{j}:\mathbb{R}\rightarrow\mathbb{R}$ in layer $j$ for each $j\in\{1,\ldots,\ell\}$, is the function

where, for each $j\in\{1,\ldots,\ell\}$, the matrix of weight $W_{j}$ (identified above with the linear map $x\mapsto W_{j}x$) has size $n_{j}\times n_{j-1}$, and $f_{j}^{\times n_{j}}$ maps $(x_{1},\ldots,x_{n_{j}})\in\mathbb{R}^{n_{j}}$ to $(f_{j}(x_{1}),\ldots,f_{j}(x_{n_{j}}))$. A typical point of view is that the weight matrices $W_{j}$ are the parameters of the neural network, while the number of layers, the number of neurons in individual layers and the activation functions are fixed for the given problem.

Once the class of neural networks is fixed, the regression problem is then to find a neural network $A$ as above which is a good approximation to a solution of the minimization problem

Solving this minimization problem is in general very difficult. A key observation is that an algorithm based on the gradient descent (or its variants, especially the stochastic gradient descent) called backpropagation has proven effective in practice. The process of (numerically, approximately) solving the minimization problem is referred to as the training of the neural network.

Under advantageous conditions, we can then interpolate between the values given in the sample set $S$ by means of a trained neural network $A$, thus obtaining a (hopefully good) approximation to $F$ on its domain $U$. However, unlike in the linear case, the obtained neural network $A$ is a black-box estimator in that it reveals no particular structure of the function $F$. It is therefore nontrivial to use $A$ for anything more than interpolating the sample data.

Given a set of sample points $\{(x,F(x)):x\in X\}$ with $X\subseteq U$ finite (and well-distributed over $U$, for instance, $\varepsilon$-dense for an $\varepsilon$ that is small enough for the scale of the problem), estimating the partial derivatives of $F$ is a difficult problem. Here are some reasons why:

• •

  The sample set $X$ is random and thus, given $x=(x_{1},\ldots,x_{d})\in X$, we cannot assume that it contains points of the form $(x_{1}+\Delta x_{1},x_{2},\ldots,x_{d})$, etc, aligned with $x$ in one of the cardinal directions. Consequently, estimating partial derivatives involves a change of basis at each point, incurring a lot of computations and numerical errors.

• •

  Estimating a partial derivative is difficult in general. One way to see the statistical issue is as follows. If $x\in X$ and $x+h\in X$ then $\frac{F(x+h)-F(x)}{h}$ can be used as an approximation for $J(x)$, but any error on $F(x)$ and $F(x+h)$, even if the error is small, is amplified upon dividing by the small quantity $h$. More formally, since we really only know $F(x)+\epsilon_{x}$ and $F(x+h)+\epsilon_{x+h}$, then $\frac{\epsilon_{x+h}-\epsilon_{x}}{h}$ will typically be large for small $h$. If the noise variables $\epsilon_{x}$ and $\epsilon_{x+h}$ have variance $\sigma^{2}$, then the variance of $\frac{\epsilon_{x+h}-\epsilon_{x}}{h}$ is $\frac{2\sigma^{2}}{h}$, which is large when $h$ is small.

• •

  Even if we could surmount the above two issues, we would still need some regression technique, such as neural networks, to interpolate the values of partial derivatives at inputs not contained in $X$.

Let us first illustrate the main idea of the algorithm in the simplest case $c=d=1$. Suppose that we wish to train a neural network $\widehat{J}$ for approximating the derivative of $F:\mathbb{R}\to\mathbb{R}$ from a finite sample set $\{(x,F(x)):x\in X\}$. Suppose that in the process of training $\widehat{J}$ we encounter a sample point $(a,F(a))$. Let $(b,F(b))$ be another sample point. If $b$ is close enough to $a$ then $F(b)$ is approximately equal to

Since $\widehat{J}$ is being trained to approximate the unknown derivative $F^{\prime}$, we naturally consider the known quantity

instead and compare it to $F(b)$, see Figure 2.

We use the related normalized quantity

as a contribution to the loss function.

In the general case of estimating the Jacobian matrix of $F:\mathbb{R}^{d}\to\mathbb{R}^{c}$, the only differences from the $\mathbb{R}\to\mathbb{R}$ case are:

• •

  The neural network $\widehat{J}$ takes a point in $\mathbb{R}^{d}$ as an input and returns a real $c\times d$ matrix.

• •

  While training $\widehat{J}$ at $a\in X$, we do not use a randomly chosen sample point $b\in X$ close to $b$ but rather the $k_{max}$ nearest neighbors of $a$ within radius $r_{max}$, where $k_{max}$ and $r_{max}$ are parameters of the algorithm. The nearest neighbors are precalculated and processed in batches.

• •

  The loss function (IV.1) is replaced with its multivariate analog, that is, for every data point $(a,F(a))$ and its neighbor $(b,F(b))$, the contribution to the loss is

  $$\frac{\left\|{F(b)-(F(a)+\widehat{J}(a)(b-a))}\right\|_{\mathbb{R}^{c}}^{2}}{\left\|{b-a}\right\|_{\mathbb{R}^{d}}^{2}}\text{.}$$

  (IV.2)

A pseudo-code for the Jacobian matrix estimator can be found in Algorithm 1.

The algorithm uses standard methods of stochastic gradient descent with backpropagation but it relies on a novel loss function and data preparation which utilizes a nearest neighbor search. An implementation of the algorithm in TensorFlow 2.5.0 [1] can be found in the appendix.

The sample set is represented as two lists $X\subseteq\mathbb{R}^{d}$ and $Y\subseteq\mathbb{R}^{c}$ of the same length $N$. For every index $j$, we think of the pair $(X[j],Y[j])$ as a sample point $(x,F(x))$, where $F$ is some unknown differentiable function $F:U\to\mathbb{R}^{c}$, $U\subseteq\mathbb{R}^{d}$.

In Algorithm 1, the initial neural network $\widehat{J}$ has $d$ cells in its input layer and $c\times d$ cells in its output layer. The function

returns all indices $0\leq j<N$ such that $X[j]\neq X[i]$ is among the $k_{max}$ nearest neighbors of $X[i]$ in $X$ and $\left\|{X[i]-X[j]}\right\|_{\mathbb{R}_{d}}<r_{max}$.

Let us point out additional features of Algorithm 1 that were used in our implementation and that are not necessarily captured in the pseudo-code.

• •

  In order not to include many nearby points in one batch, we randomly permute the entries of $D$ at the end of the data preparation stage.

• •

  Even if $|X|=N=|Y|$ and the parameters $k_{max}$ and $r_{max}$ are fixed, then number of training points in each run of the algorithm depends on $X$ (since we do not know in advance how many nearest neighbors of a given point will satisfy the constraint on $r_{max}$). After the training data $D$ is prepared, we adjust the batch size so that it divides the size of $D$ evenly. Note that the training set therefore has cardinality $|D|$ satisfying $N\leq|D|\leq Nk_{max}$.

• •

  We allow to set the value of $r_{max}$ to infinity so that every nearest neighbor set has cardinality $k_{max}$ (if $k_{max}<N$).

• •

  The effect of the parameter $k_{max}$ can be suppressed by setting its value to at least $N$.

Note that it is possible for both $(i,j)$ and $(j,i)$ to occur in the training set $D$. If $(i,j)\in D$, which means that $X[j]$ is one of the $k_{max}$ nearest neighbors of $X[i]$ and within radius $r_{max}$ of $X[i]$, it does not necessarily follow that $(j,i)\in D$, too. With each $(i,j)\in D$, the neural network $\widehat{J}$ will be trained using a linear approximation centered at $X[i]$.

Having trained the neural network $\widehat{J}$ by Algorithm 1, we can obtain an approximate value of the Jacobian matrix $J$ of the unknown function $F:U\to\mathbb{R}^{c}$ by computing $\widehat{J}(x)$ for $x\in U$.

Note that we can also obtain an approximate value of the unknown function $F$ itself at $x\in U$ by locating a nearby point $y$ in the sample set and computing $F(y)+\widehat{J}(y)(x-y)$.

In all examples of Subsections V-A and V-B we set the parameters of the algorithm as

• •

  $r_{max}=0.5$,

• •

  $k_{max}=30$,

and of the neural network as

• •

  a simple forward ANN,

• •

  with $4$ hidden layers consisting of $100$, $100$, $50$ and $20$ neurons,

• •

  using the swish activation function for all cells,

• •

  and trained with stochastic gradient descent in $50$ epochs and batch size of $50$.

Consider the function $F_{0}$ of Table I. This function is not known to Algorithm 1 for the purposes of training the Jacobian matrix estimator $\widehat{J}$, nevertheless it is known to us and we can therefore visualize it, calculate its Jacobian matrix $J$ by standard symbolic differentiation, and visualize the Jacobian matrix as well.

The graph of $F_{0}$ is given in Figure 3(a). Its Jacobian matrix $J$ is a function $\mathbb{R}^{2}\to\mathbb{R}^{2}$ and we plot it as a vector field on a regular grid of $20\times 20$ points contained in $(-2,2)^{2}$. This is done in Figure 3(c), where the vectors are automatically scaled to improve legibility, and in Figure 3(e), where the vectors are to scale.

The input of the algorithm is a cloud of $N=10^{6}$ sample points $\{(x,F(x)):x\in X\}$, where the set $X\subseteq(-2,2)^{2}$ is generated randomly. Such a cloud is visualized in Figure 3(b), except that only $10^{3}$ points are plotted in the figure to improve legibility.

The resulting Jacobian matrix estimator $\widehat{J}$ is visualized as a vector field in Figure 3(d) automatically scaled, and in Figure 3(f) to scale.

In order to compare the near-identical vector fields of $J$ and $\widehat{J}$, we offer Figure 4 in which we plot the vector field $\widehat{J}-J$ to scale.

It can be seen plainly from Figure 4 that $\widehat{J}$ is a good estimator of $J$ for the function $F_{0}$. In order to quantify this fact for $F_{0}$ and for the other testing functions $F:\mathbb{R}^{d}\to\mathbb{R}^{c}$ from Table I, we consider two validation methods.

For the first validation method, suppose that the Jacobian matrix estimator is trained for a function $F:\mathbb{R}^{d}\to\mathbb{R}^{c}$ that is not known to the algorithm but that is known to us. We can therefore compare the estimated Jacobian matrix $\widehat{J}$ with the actual Jacobian matrix $J$ of $F$ calculated by standard symbolic differentiation.

In more detail, let $S$ be a very large, randomly generated subset of the domain of $F$. (We used $|S|=10^{6}$ throughout.) For $\delta\geq 0$, let

where $\left\|{J(x)}\right\|_{\mathbb{R}^{cd}}$ is the Frobenius norm of the matrix $J(x)$, that is, the Euclidean norm of the vector in $\mathbb{R}^{cd}$ obtained by concatenating the rows of $J(x)$. Consider the error estimate

that averages pointwise over $S_{\delta}$ the size of the error $\widehat{J}(x)-J(x)$ relative to the size of $J(x)$. We will refer to $E_{\delta}$ as the average relative error of $\widehat{J}$ over $S$ and report it in percents.

For the second validation method, suppose that we would like to train the Jacobian matrix estimator on a sample set $\{(x,F(x)):x\in X\}$ for a function $F:\mathbb{R}^{d}\to\mathbb{R}^{c}$ that is not known to us. We split $X$ into two disjoint subsets $X_{T}$ and $X_{V}$, a larger training set $X_{T}$ and a smaller validation set $X_{V}$. (We used $|X_{V}|=10^{4}$ except as otherwise noted.) We then train $\widehat{J}$ on the training set $\{(x,F(x)):x\in X_{T}\}$ and we calculate the error on the validation set $X_{V}$ as follows.

Let $X^{*}_{V}$ be the set of all pairs $(a,b)\in X_{V}\times X_{V}$ such that $b$ is a near neighbor of $a$ obtained by the same near neighbor routine that has been employed in Algorithm 1. For $\delta>0$, let

and note that $F(b)$ is known from the sample set. Instead of comparing $\widehat{J}$ to the unknown Jacobian matrix $J$, consider the error estimate

based on linear approximations. The purpose of the parameter $\delta$ is the same as in the first error estimate (V.1). We again report $E^{*}_{\delta}$ in percents.

Table II summarizes the error estimates for all testing functions from Table I.

The effect of $\delta$ is most visible for functions whose Jacobian matrix is frequently close to $0$. Not surprisingly, the error estimate $E^{*}_{\delta}$ tends to be better than $E_{\delta}$ since both the training of $\widehat{J}$ and the error estimate $E^{*}_{\delta}$ are based on linear approximations, albeit at different pairs of points. Note that the error gets worse as the dimension $d$ of the domain increases.

The rather large error $E^{*}_{0.01}$ for $F_{11}$ and $F_{12}$ is a consequence of the large volume of the domains and small validation set $X_{V}$. Increasing the size of $X_{V}$ from $10^{4}$ (the default) to $10^{5}$ for $F_{11}$ (resp. $F_{12}$) with $N=10^{6}$ improves $E^{*}_{0.01}$ from $10.3$ (resp. $8.47$) to $3.92$ (resp. $3.53$).

We conclude this section with somewhat informal comments on the effects of varying the parameters and on the limitations of Algorithm 1. A more thorough discussion will be reported elsewhere.