SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Amir Akbarnejad
Department of Computing Science
Edmonton, AB, Canada
ah8@ualberta.ca
Gilbert Bigras
Department of Laboratory Medicine and Pathology
Gilbert.Bigras@albertaprecisionlabs.ca
Nilanjan Ray
nray1@ualberta.ca
The analogy between Gaussian processes (GPs) and deep artificial neural networks (ANNs) has received a lot of interest, and has shown promise to unbox the blackbox of deep ANNs. Existing theoretical works put strict assumptions on the ANN (e.g. requiring all intermediate layers to be wide, or using specific activation functions). Accommodating those theoretical assumptions is hard in recent deep architectures, and those theoretical conditions need refinement as new deep architectures emerge. In this paper we derive an evidence lower-bound that encourages the GP's posterior to match the ANN's output without any requirement on the ANN. Using our method we find out that on 5 datasets, only a subset of those theoretical assumptions are sufficient. Indeed, in our experiments we used a normal ResNet-18 or feed-forward backbone with a single wide layer in the end. One limitation of training GPs is the lack of scalability with respect to the number of inducing points. We use novel computational techniques that allow us to train GPs with hundreds of thousands of inducing points and with GPU acceleration. As shown in our experiments, doing so has been essential to get a close match between the GPs and the ANNs on 5 datasets. We implement our method as a publicly available tool called GPEX: https://github.com/amirakbarnejad/gpex. On 5 datasets (4 image datasets, and 1 biological dataset) and ANNs with 2 types of functionality (classifier or attention-mechanism) we were able to find GPs whose outputs closely match those of the corresponding ANNs. After matching the GPs to the ANNs, we used the GPs' kernel functions to explain the ANNs' decisions. We provide more than 200 explanations (around 30 explanations in the paper and the rest in the supplementary) which are highly interpretable by humans and show the ability of the obtained GPs to unbox the ANNs' decisions.
Artificial neural networks (ANNs) are widely adopted in machine learning. Despite their benefits, ANNs are known to be black-box to humans, meaning that their inner mechanism for making predictions is not necessarily interpretable/explainable to humans. ANN's black-box property impedes its deployment in safety-critical applications like medical imaging or autonomous driving, and makes them hard-to-troubleshoot for machine learning researchers.
Attribution-based explanation methods like LIME[31], SHAP[21] and most gradient-based explanation methods like DeepLIFT [5] presume a linear surrogate model. Given a test instance  $x_{test}$ , this simpler surrogate model is encouraged to have the same output "locally" around  $x_{test}$ . Because of this "local assumptions", explanations from these methods might be unreliable, and can be easily manipulated by an adversary model [11][28]. Moreover, these models may produce discordant explanations for a fixed model and test instance [16].
Considering Gaussian processes (GPs) [26] as the explainer model is beneficial, because: 1. Gaussian processes are highly interpretable. 2. Researchers have long known that GP's posterior has the potential to match an ANN's output "globally". More precisely, given an ANN and some requirements on it [24][8], there might exist a GP whose posterior matches the ANN's output all over the input-space  $\mathbb{X}$  (as opposed to the local explanation models for which the match happens only locally around a test instance  $x_{test} \in \mathbb{X}$ ). Not many explainer models can globally match the ANN's output. Among gradient-based methods, with the best of our knowledge only Integrated Gradients [33] has a weak sense of ANN's global behaviour over the input space. Having some conditions on an ANN, represent point selection [37] finds a "globally faithful" explainer model that, similar to GPs, works with a kernel function. As we will elaborate upon in Sec. 4.3 and Sec. S6 of the supplementary, the GP's kernel that we find in this paper is superior due to a technical point in the formulation of representer point selection [37]. All in all, using GPs to explain ANNs is quite promising and has advantages over other approaches to explain ANNs.
- Theoretical results on ANN-GP analogy impose some restrictions on ANNs under which the ANN will be equivalent to a GP. These conditions are too restrictive for recently used deep architectures. Moreover, those theoretical conditions need refinement as new deep architectures emerge. In this paper we derive an ELBO for training GPs which encourages GP's posterior to match ANN's output. Our formulation and method doesn't impose any restriction on the ANN and the method used to train it.
- Using our method, we empirically show that on 5 datasets (4 image datasets, and 1 biological dataset) and ANNs with 2 types of functionality (classifier or attention-mechanism) the ANN needs to fulfill only a subset of those theoretical conditions. Indeed, in our experiments we used a normal ResNet-18 or feed-forward backbone with a single wide layer in the end.
- Scalability is a major issue in training GPs. To address this issue, we adopted computational techniques recently used for fast spectral clustering [13] as well as a novel method to learn the GPs using mini-batches of inducing points and training instances. These computational techniques allow us to train GPs with hundreds of thousands of inducing points. According to our analysis, increasing the inducing points has been essential to get a good match between the trained GPs and ANNs. Indeed, without many inducing points GPs posterior cannot be a complex function (a function with many ups and downs [36]) and fails to match ANNs' output.
- With the best of our knowledge, our work is the first method that performs knowledge distillation between GPs and ANNs.
- We implement our method as a public python library called GPEX (Gaussian Processes for EXplaining ANNs). GPEX takes in an arbitrary PyTorch module, and replaces any ANN submodule of choice by GPs. Our package makes use of GPU-acceleration, and enables effortless application of GPs without getting users involved in details of the inference procedure. GPEX can be used by machine learning researchers to interpret/troubleshoot their artificial neural networks. Moreover, GPEX can be used by researchers working on the theoretical side of ANN-GP analogy to empirically test their hypotheses.
In this article the function  $g(\cdot)$  always denotes an ANN. The kernel of a Gaussian process is denoted by the double-input function  $\mathcal{K}(\cdot, \cdot)$ . We assume the kernel similarity between two instances  $x_i$  and  $x_j$  is equal to  $f(\boldsymbol{x}_i)^T f(\boldsymbol{x}_j)$ , where  $f(\cdot)$  maps the input-space to the kernel space. In this paper  $\boldsymbol{u}$  (resp.  $v$ ) denotes a vector in the kernel-space (resp. the posterior mean) of a GP. In some sense  $\boldsymbol{u}$  and  $v$  denote
Figure 1: a) A general feed-forward pipeline, with an ANN sub-module to be explained by GPEX. b) Typical behaviour of Gaussian process posterior given a set of observed values.
the input and the output of a GP, respectively. We have that:  $\mathcal{K}(\pmb{x}_i,\pmb{x}_j) = f(\pmb{x}_i)^Tf(\pmb{x}_j) = \pmb{u}_i^T\pmb{u}_j$ . The number of GPs is equal to the number of the outputs from the ANN. In other words, we consider one GP per scalar output-head from the ANN. We use index  $\ell$  to specify the  $\ell$ -th GP as follows:  $\mathcal{K}_{\ell}(\pmb{x}_i,\pmb{x}_j) = f_{\ell}(\pmb{x}_i)^Tf_{\ell}(\pmb{x}_j) = \pmb{u}_{i}^{(\ell)T}\pmb{u}_{j}^{(\ell)}$ . We parameterize the  $\ell$ -th GP by a set of  $M$  inducing points  $\{(\tilde{\pmb{u}}_m^{(\ell)},\tilde{v}_m^{(\ell)})\}_{m=1}^M$ . The tilde in  $(\tilde{\pmb{u}}_m^{(\ell)},\tilde{v}_m^{(\ell)})$  indicates that  $\tilde{\pmb{u}}$  is one of the  $M$  inducing points in the kernel space. However,  $\pmb{u}$  (without tilde) can be an arbitrary point in the continuous kernel-space.
2.2 The Proposed Framework
To make our framework as general as possible, we consider a general feed-forward pipeline that contains an ANN as a submodule. In Fig. 1a the bigger square illustrates the general module. The input-output of the general pipeline are denoted in Fig. 1a by  $\mathcal{X}$  and  $\mathcal{Y}$ . The general pipeline has at least one ANN submodule to be explained by GPEX. Fig. 1a illustrates this ANN by the small blue rectangle within the general pipeline. The input-output of the ANN are denoted in Fig. 1a by  $x$  and  $v$ . Note that  $\mathcal{X}$  and  $\mathcal{Y}$  can be anything, including without any limitation, a set of vectors, labels, and meta-information. However, input-output of the ANN (i.e.  $x$  and  $y$ ) are required to be in tensor format. The exact requirements are provided in the online documentation for GPEX. Moreover, the general module can have other arbitrary submodules, which are depicted by the blue clouds. The relations between the submodules, as illustrated by the dotted-lines in Fig. 1a, can also be quite general. Our probabilistic formulation only needs access to the conditional distributions  $p(\boldsymbol{x}|\mathcal{X})$  and  $p(\mathcal{Y}|\boldsymbol{x},\mathcal{X})$ . Similarly, the proposed GPEX is completely agnostic about the general pipeline and it only requires the ANN's input-output to be in the tensor format. Given a PyTorch module, the proposed GPEX tool automatically grabs the distributions  $p(\boldsymbol{x}|\mathcal{X})$  and  $p(\mathcal{Y}|\boldsymbol{x},\mathcal{X})$  from the main module it is given.
The inducing points  $\{\tilde{\pmb{u}}_m^{(\ell)},\tilde{v}_m^\ell \}_{m = 1}^M$  parameterize the  $\ell$  -th GP. Note that  $\tilde{\pmb{u}}_m^{(\ell)} = f_\ell (\tilde{\pmb{x}}_m)$  . A feature point like  $\pmb{x}$  is first mapped to the kernel-space as  $\pmb{u}^{(\ell)} = f_{\ell}(\pmb {x})$  . Note that the kernel functions  $\{f_{\ell}(.)\}_{\ell = 1}^{L}$  are implemented as separate neural networks, or for the sake of efficiency as a single neural network backbone with  $L$  different heads. Afterwards, the GP's posterior on  $\pmb{x}$  depends on the kernel similarities between  $\pmb{u}^{(\ell)}$  and the inducing points  $\{\tilde{\pmb{u}}_m^{(\ell)}\}_{m = 1}^M$  . More precisely, the posterior of the  $\ell$  -th GP on  $\pmb{x}$  is a random variable  $v^{(\ell)}$  whose distribution is as follows [26]:
$$
p (v ^ {(\ell)} | \boldsymbol {u} ^ {(\ell)}, \tilde {\boldsymbol {u}} _ {1: M} ^ {(\ell)}, \tilde {v} _ {1: M} ^ {(\ell)}) = \mathcal {N} \left(v ^ {(\ell)}; \mu_ {v} (\boldsymbol {u} ^ {(\ell)}, \tilde {\boldsymbol {u}} _ {1: M} ^ {(\ell)}, \tilde {v} _ {1: M} ^ {(\ell)}), c o v _ {v} (\boldsymbol {u} ^ {(\ell)}, \tilde {\boldsymbol {u}} _ {1: M} ^ {(\ell)}, \tilde {v} _ {1: M} ^ {(\ell)})\right), \tag {1}
$$
where  $\mu_v(.,.,.)$  and  $cov_v(.,.,.)$  are the mean and covariance of a GP's posterior computed as:
$$
\mu_ {v} \left(\boldsymbol {u} ^ {(\ell)}, \tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}, \tilde {v} _ {1: M} ^ {(\ell)}\right) = \mathcal {K} \left(\boldsymbol {u} ^ {(\ell)}, \tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}\right) \left[ \mathcal {K} \left(\tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}, \tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}\right) + \sigma_ {g p} ^ {2} \mathbf {I} _ {M \times M} \right] ^ {- 1} \tilde {v} _ {1: M} ^ {(\ell)} \tag {2}
$$
$$
\begin{array}{l} c o v _ {v} (\boldsymbol {u} ^ {(\ell)}, \tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}, \tilde {v} _ {1: M} ^ {(\ell)}) = \mathcal {K} (\boldsymbol {u} ^ {(\ell)}, \boldsymbol {u} ^ {(\ell)}) - \\ \mathcal {K} \left(\boldsymbol {u} ^ {(\ell)}, \tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}\right) \times \left[ \mathcal {K} \left(\tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}, \tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}\right) + \sigma_ {g p} ^ {2} \mathbf {I} _ {M \times M} \right] ^ {- 1} \mathcal {K} \left(\tilde {\mathbf {u}} _ {1: M} ^ {(\ell)}, \boldsymbol {u} ^ {(\ell)}\right). \tag {3} \\ \end{array}
$$
As the variables  $\{v_m^{(\ell)}\}_{m=1}^M$  and  $v$  are latent or hidden, we train the model parameters by optimizing a variational lower-bound. We consider the following variational distributions:
$$
q _ {1} (v ^ {(\ell)} \mid \boldsymbol {x}) = \mathcal {N} \big (v ^ {(\ell)} ; g _ {\ell} (\boldsymbol {x}), \sigma_ {g} ^ {2} \big), \quad q _ {2} \big (\tilde {v} _ {m} ^ {(\ell)} \big) = \mathcal {N} \big (\tilde {v} _ {m} ^ {(\ell)} ; \varphi_ {m} ^ {(\ell)}, \sigma_ {\varphi} ^ {2} \big). \tag {4}
$$
In Eq. 4, the function  $g_{\ell}(.)$  is the  $\ell$ -th output from the ANN. Note that as the set of hidden variables  $\{\tilde{v}_m^{(\ell)}\}_{m=1}^M$  is finite, we have parameterized their variational distribution by a finite set of numbers  $\{\varphi_m^{(\ell)}\}_{m=1}^M$ . However, as the variables  $x$  can vary arbitrarily in the feature space, the variable  $u^{(\ell)}$  varies arbitrarily in the kernel space. Therefore, the set of values  $v^{(\ell)}$  may be infinite. Accordingly, the variational distribution for  $v^{(\ell)}$  is conditioned on  $x$  and is parameterized by the ANN  $g(.)$ .
2.3 The Derived Evidence Lower-Bound (ELBO)
Due to space limitation, the derivation of the lower-bound is moved to Sec. S1 of the supplementary material. In this section we only introduce the derived ELBO and discuss how it relates the GP, the ANN and the training cost of the main module in an intuitive way. The ELBO terms containing the GP parameters (i.e. the parameters of the kernel function  $f(\cdot)$ ) is denoted by  $\mathcal{L}_{gp}$ . According to Eq. S9 of the supplementary material  $\mathcal{L}_{gp}$  is as follows:
$$
\begin{array}{l} \mathcal {L} _ {g p} = - \frac {1}{2} \mathbb {E} _ {\sim q} \left[ \sum_ {\ell = 1} ^ {L} \frac {\left(\mu_ {v} (\boldsymbol {u} ^ {(\ell)} , \tilde {\mathbf {u}} _ {1 : M} ^ {(\ell)} , \tilde {v} _ {1 : M} ^ {(\ell)}) - g _ {\ell} (\boldsymbol {x})\right) ^ {2} + \sigma_ {g} ^ {2}}{c o v _ {v} (\boldsymbol {u} ^ {(\ell)} , \tilde {\mathbf {u}} _ {1 : M} ^ {(\ell)} , \tilde {v} _ {1 : M} ^ {(\ell)})} \right] \\ - \frac {1}{2} \mathbb {E} _ {\sim q} \left[ \sum_ {\ell = 1} ^ {L} \log \left(\frac {\operatorname {c o v} _ {v} \left(\boldsymbol {u} ^ {(\ell)} , \tilde {\mathbf {u}} _ {1 : M} ^ {(\ell)} , \tilde {v} _ {1: M} ^ {(\ell)}\right)}{\sigma_ {g} ^ {2}}\right) \right] + (\text {c o n s t .}), \tag {5} \\ \end{array}
$$
where  $q(\cdot)$  is the variational distribution that factorizes to the  $q_{1}(.)$  and  $q_{2}(.)$  distributions defined in Eq. 4. In the first term of Eq. 5, the numerator encourages the GP and the ANN to have the same output. More precisely, for a feature point  $x$  we can compute the corresponding point in the kernel space as  $\boldsymbol{u}^{(\ell)} = f_{\ell}(\boldsymbol{x})$  and then compute the GP's posterior mean based on kernel similarities between  $\boldsymbol{u}$  and the inducing points to get the GP's mean  $\mu_v$ . In Eq. 5 the GP's mean  $\mu_v$  is encouraged to match the ANN's output  $g_{\ell}(\boldsymbol{x})$ . In Eq. 5, because of the denominator of the first term, the ANN-GP similarity is not encouraged uniformly over the feature-space. Wherever the GP's uncertainty is low, the term  $cov_v(\boldsymbol{u}^{(\ell)}, \tilde{\mathbf{u}}_{1:M}^{(\ell)}, \tilde{v}_{1:M}^{(\ell)})$  in the denominator becomes small. Therefore, the GP's mean is highly encouraged to match the ANN's output. On the other hand, in regions where the GP's uncertainty is high, the GP-ANN analogy is less encouraged. This formulation is quite intuitive according to the behaviour of Gaussian processes. Fig. 1b illustrates the posterior of a GP with radial-basis kernel for a given set of observations. In regions like  $[3,\infty)$  and  $(-\infty, -4]$  there are no nearby observed data. Therefore, in these regions the GP is highly uncertain and the blue uncertainty margin is thick in such regions. Intuitively, our derived ELBO in Eq. 5 encourages the GP-ANN analogy only when GP's uncertainty is low and gives less importance to regions similar to  $[3,\infty)$  and  $(-\infty, -4]$  in Fig. 1b. Note that this formulation makes no difference for the ANN as ANNs are known to be global approximators. However, this formulation makes a difference when training the GP, because the GP is not required to match the ANN in regions where there are no similar training instances. The ELBO terms containing the ANN parameters is denoted by  $\mathcal{L}_{ann}$ . According to Sec. S1.2 of the supplementary material,  $\mathcal{L}_{ann}$  is as follows:
$$
\mathcal {L} _ {a n n} = - \frac {1}{2} \mathbb {E} _ {\sim q} \left[ \sum_ {\ell = 1} ^ {L} \frac {\left(\mu_ {v} \left(\boldsymbol {u} ^ {(\ell)} , \tilde {\mathbf {u}} _ {1 : M} ^ {(\ell)} , \tilde {v} _ {1 : M} ^ {(\ell)}\right) - g _ {\ell} (\boldsymbol {x})\right) ^ {2}}{\operatorname {c o v} _ {v} \left(\boldsymbol {u} ^ {(\ell)} , \tilde {\mathbf {u}} _ {1 : M} ^ {(\ell)} , \tilde {v} _ {1 : M} ^ {(\ell)}\right)} \right] + \mathbb {E} _ {\sim q} \left[ \log p (\mathcal {Y} | \boldsymbol {y}, \mathcal {X}) \right]. \tag {6}
$$
In the above objective the first term encourages the ANN to have the same output as the GP. Similar to the objective of Eq. 5, the denominator of the first term gives more weight to ANN-GP analogy when GP's uncertainty is low. In the right-hand-side of Eq. 6, the second term is the likelihood of the pipeline's output(s), i.e.  $\mathcal{V}$  in Fig. 1a. This term can be, e.g., the cross-entropy loss when  $\mathcal{V}$  contains class scores in a classification problem, or the mean-squared error when  $\mathcal{V}$  is the predicted value for a regression problem, or a combination of those costs in a multi-task setting.
We consider a separate Gaussian process for each output head of an ANN. In other words, given an ANN we have as many GPs as the number of the ANN's output heads. To explain an ANN, we find the explainer GPs by optimizing the objective in Eq. 5 w.r.t. to the kernel mappings  $\{f_{\ell}(.)\}_{\ell = 1}^{L}$ . To do so, we need to have  $\mu_v$  which in turn means we need to have all kernel-space representations
Algorithm 1 Method Optim_KernMappings
Input: Input instance  $\pmb{x}$  and inducing instance  $\tilde{\pmb{x}}$  , list of matrices U, list of vectors V.   
Output: Kernel-space mappings  $[f_1(.),\dots,f_L(.)]$  . Initialization : loss  $\leftarrow 0$  1:  $\mu ,cov\gets$  forward_GP(x,  $\tilde{\pmb{x}}$  ,U,V)//feed x to GPs,"forward_GP" is Alg.S1 in supplementary. 2:  $\mu_{ann}\gets g(\mathbf{x})$  //feed x to ANN. 3: for  $\ell = 1$  to  $L$  do 4: loss  $\leftarrow$  loss  $+\frac{(\mu[\ell] - \mu_{ann}[\ell])^2 + \sigma_g^2}{cov[\ell]} +\log (cov[\ell]).$  //Eq.5. 5: end for 6:  $\delta \gets \frac{\partial\textit{loss}}{\partial\textit{params}([f_1(.),\ldots,f_L(.)])}$  //the gradient of loss. 7: params([f1,...,fL])  $\leftarrow$  params([f1,...,fL])-lr×δ //update the parameters. 8: lr← updated learning rate 9: return [f1(.),...,fL(.)]
Algorithm 2 Method Explain ANN
Input: Training dataset ds_train, and the inducing dataset ds_inducing.  
Output: Updated kernel-space mappings  $[f_1(.),\dots,f_L(.)]$  , and the other GP parameters U and V. Initialization : U, V  $\leftarrow$  Init_GPparams(ds_inducing) //Alg.S3 in supplementary.  
1: for iter = 1 to max_iter do  
2:  $\pmb{x}\gets$  randselect(ds_train).  
3:  $\tilde{\pmb{x}}\gets$  randselect(ds_inducing)  
4:  $[f_1(.),\dots,f_L(.)]\gets$  Optim_KernMapings(x,  $\tilde{x}$  ,U,V).  
5:  $\tilde{\pmb{x}}\gets$  randselect(ds_inducing).  
6: for  $\ell = 1$  to  $L$  do  
7: //update kernel-space representations.  
8:  $U[\ell][\tilde{x}.index]\gets f_{\ell}(\tilde{x})$   
9: end for  
10: end for  
11: return  $[f_1(.),\dots,f_L(.)]$  , U, V
$\{\tilde{u}_m^{(\ell)}\}_{m = 1}^M$  . However, it is computationally prohibitive to feed thousands of inducing instances to the kernel mappings as  $\tilde{\pmb{u}}_m^{(\ell)} = f_\ell (\tilde{\pmb{x}}_m)$  for  $m\in \{1,2,\dots,M\}$  in each gradient-descent iteration. On the other hand, as the kernel-space mappings  $\{f_{\ell}(.)\}_{\ell = 1}^{L}$  keep changing during training, we need to somehow track how the inducing points  $\{\tilde{u}_m^{(\ell)}\}_{m = 1}^M$  change during training. To this end, we put the kernel-space representations of the inducing points in matrices denoted by  $\mathbf{U}$  . During training, these matrices are repeatedly updated by feeding mini-batches of inducing instances to the kernel-mappings.
Alg. 2 optimizes the objective of Eq. 5 w.r.t. the kernel mappings  $\{f_{\ell}(.)\}_{\ell = 1}^{L}$ . First, a single training instance  $x$  and a single inducing point  $\tilde{x}$  are selected (line 2-3). Afterwards, the procedure of Alg. 1 is called to update the kernel mappings (line 4 of Alg. 2). To update the kernel-mappings, the GP posterior is computed via the matrices  $\mathbf{U}$  (line 1 of Alg.1). The "forward_GP" procedure (called in line 1 of Alg. 1) is provided in Alg. S1 of the supplementary, and uses the matrices  $\mathbf{U}$  to compute GP's posterior. Only the rows of  $\mathbf{U}$  that correspond to the selected inducing point  $\tilde{x}$  are computed using the kernel-mappings, so that the gradient w.r.t. the kernel-mappings can be computed in the backward pass (lines 5-6 of Alg. S1 in the supplementary). Finally, the matrices  $\mathbf{U}$  are updated (lines 6-8 of Alg. 2). Due to the lack of space, the routines "forward_GP" and "Init_GPparams" and more details are moved to Sec. S2 of the supplementary. Of course instead of a single training/inducing instance, we used a mini-batch of multiple training/inducing instances.
One difficulty of training GPs is the matrix inversion of Eqs.2 and 3, which has  $\mathcal{O}(M^3)$  complexity using standard matrix inversion methods. To address this issue, we adopted computational techniques recently used for fast spectral clustering [13]. Let  $\mathbf{A}$  be an arbitrary  $M\times D$  matrix where  $M >> D$ . Moreover, let  $\mathbf{b}$  be a  $M$ -dimensional vector and let  $\sigma$  be a scalar. The computational techniques [13] allow us to efficiently compute:  $(\mathbf{AA}^T + \sigma^2\mathbf{I}_{M\times M})^{-1}\mathbf{b}$ . (Note the similar terms in the right hand side of Eqs. 2 and 3.) The idea is that  $\mathbf{AA}^T$  is of rank  $D$ . Therefore, from linear algebra it follows
that  $(\mathbf{A}\mathbf{A}^T + \sigma^2\mathbf{I}_{M\times M})^{-1}$  has  $M - D$  eigen-values all of which are equal to  $\sigma^{-2}$ . Therefore, in the space of those eigen-vectors, the transformation on  $\pmb{b}$  is simply a scaling by  $\sigma^{-2}$ . The details and more computational techniques are provided in Sec. S2.1 of the supplementary. These computational techniques allow us to efficiently compute the GP-posterior for hundreds of thousands of inducing points in each gradient descent iteration.
A note on the used datasets in Alg. 2: According to our analysis of Sec. S7 in the supplementary material, "ds_inducing" should be as large as possible so the GP posteriors can be flexible enough to match the ANNs. Therefore a good practice is to include all training instances (without data augmentation) in "ds_inducing". By doing so, the following issue arises. An instance from "ds_train" like  $\pmb{x}$  is an augmented version of an inducing instance  $\tilde{\pmb{x}}$ . Because  $\pmb{x}$  and  $\tilde{\pmb{x}}$  are close, their kernel-space representations  $f(\pmb{x})$  and  $f(\tilde{\pmb{x}})$  also become close regardless of parameters of  $f(.)$ . Consequently, regardless of  $f(.)$ , GP's posterior mean will be roughly equal for both  $\pmb{x}$  and  $\tilde{\pmb{x}}$ . Indeed, in this case Alg. 2 fails to find the kernel mappings  $\{f_{\ell}(.)\}_{\ell=1}^{L}$ . To avoid this issue, we sample  $\pmb{x}$  in line 2 of Alg. 2 as follows:  $\pmb{x}_1$  and  $\pmb{x}_2$  are randomly selected from "ds_train", and  $\alpha \sim uniform(-1,2)$ , and  $\pmb{x} = \alpha \pmb{x}_1 + (1 - \alpha) \pmb{x}_2$ . The rest of Alg. 2 after line 2 is run as before.