Adaptive wavelet distillation from neural networks through interpretations
Authors: Wooseok Ha
Chandan Singh
Francois Lanusse
Srigokul Upadhyayula
Bin Yu
Venue: NeurIPS 2021 Poster (2021)
Keywords: Interpretability
Distillation
Wavelet
Cosmology
Cell biology

Abstract: Recent deep-learning models have achieved impressive prediction performance, but often sacrifice interpretability and computational efficiency. Interpretability is crucial in many disciplines, such as science and medicine, where models must be carefully vetted or where interpretation is the goal itself. Moreover, interpretable models are concise and often yield computational efficiency. Here, we propose adaptive wavelet distillation (AWD), a method which aims to distill information from a trained neural network into a wavelet transform. Specifically, AWD penalizes feature attributions of a neural network in the wavelet domain to learn an effective multi-resolution wavelet transform. The resulting model is highly predictive, concise, computationally efficient, and has properties (such as a multi-scale structure) which make it easy to interpret. In close collaboration with domain experts, we showcase how AWD addresses challenges in two real-world settings: cosmological parameter inference and molecular-partner prediction. In both cases, AWD yields a scientifically interpretable and concise model which gives predictive performance better than state-of-the-art neural networks. Moreover, AWD identifies predictive features that are scientifically meaningful in the context of respective domains. All code and models are released in a full-fledged package available on Github.

Adaptive wavelet distillation from neural networks through interpretations
Wooseok Ha' haysse@berkeley, edu
Chandan Singh² chandan_shnhhbberkeleey. edu
Francois Lanusse
francois. arancois.lanusseee waw.sse@cea. .fr
Srigokul Upadhyayula' sup@berkeley.cm edu
Bin Yu 12
binyu@berkeley.com edu
Statistics Department, UC Berkeley
2 EFCS Department, UC Berkeley
AIM, CEA, CNRS; Université Paris-Saclay, Université Paris Diderot, Sorbonne Paris Cité Advanced Bioimaging Center, Department of Molecular & Cell Biology, UC Berkeley
Abstract
Recent deep-learing models have achieved impressive prediction performance, but often sacrifice interpretability and computational efficiency. Interpretability is crucial in many disciplines, such science and medicine, where models must be carefully vetted where interpretation is the goal itself. Moreover, interpretable models are concise and often yield computational efficiency. Here, we propose adaptive wavelet distillation (AWD), method which aims to distill information from a trained neural network into wavelet transform. Specifically, AWD pe- nalizes feature attributions of neural network in the wavelet domain to learn an effective multi-resolution wavelet transform. The resulting model highly predic- tive, concise, computationally efficient, and has properties (such multi-scale structure) which make it easy to interpret. experts, we showcase how AWD addresses challenges in two real-world settings: cosmological parameter inference and molecular-pare AWD yields scientifically interpretable and concise model which gives predictive performance better than state-of -ft-----tthtteee the-art neural networks. Moreover, AWD identi- fies predictive features that scientifically meaningful in the context of respective domains. All code and models are released full-fledged package available on Github.
close collaboration with domain
prediction. In both cases,
Introduction
Recent advancements in deep learning have led impressive increases in predictive performance. However, the inability to interpret deep neural networks (DNNs) has led them to be characterized as black boxes. It is often critical that models are inherently interpretable [1-3], particularly in high- stakes applications such as medicine, biology, and molicy-making. In these cases, interpretations which are relevant particular domain/audience [3] can ensure that models behave reasonably, identify when models will make errors, and make the models more amenable to inspection and improvement by domain experts. Moreover, interpretable models tend be faster and more computationally efficient than large neural networks.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

One promising approach to constructing interpretable models without sacrificing prediction perfor- mance is model distillation. Model distillation [4-6] transfers the knowledge in one model (i.e., the teacher), into another model (i.e., the student), where the student model often has desirable properties, such as being more interpretable than the teacher model, Recent works have considered distilling a DNN into inherently interpretable models such as a decision tree [7-9] or global additive model [10], with some success. Here, we consider distilling a DNN into a learnable wavelet transform, which is powerful tool to describe signals both in time (spatial) and frequency domains that has found numerous successful applications physical and biomedical sciences.
Wavelet
Pre-trained DNN
transform y
X
Y
4-1
8
8
DNN attributions for wavelet coefficients
Figure I: Adaptive wavelet distillation uses attributions from a trained DNN to improve its wavelet transform, while satisfying constraints for reconstruction error and wavelet constraints. See Eq. 8 for the precise formulation of the optimization objective.
Wavelets have many properties amenable to interpretation: they can form an orthogonal basis, identify sparse representation of signal, and tile different frequencies and spatial locations (and sometimes rotations), allowing for multiresolution analysis. Most previous work has focused on hand-designed wavelets for different scenarios rather than wavelets which adapt to given data. Recent work has explored wavelets which adapt to an input data distribution, under the name optimized wavelets or adaptive wavelets [11-19]. Moreover, some work has used wavelets as part of the underlying structure of a neural network, as in wavelet networks / wavelet neural networks [20-25], or the scattering transform [26, 27]. However, none of them utilize wavelets for interpretable model distillation.
Fig outlines Adaptive Wavelet Distillation (AWD), our approach for distilling a wavelet transform from a trained DNN. A key novelty of AWD is that it uses attributions from a trained DNN to improve the learned " this incorporates information not just about the input signals, as is done in previous work, but also about the target variable and the inductive biases present in the DNN.
This paper deviates significantly from a typical NeurIPS paper. While there has been an explosion of work in "interpretable machine learning" [28], there has been very limited development and grounding of these methods in the context of a particular problem and audience. This has led to much confusion about how to develop and evaluate interpretation methods [29,30]; in fact, major part of the issue is that interpretability cannot be properly defined without the context of a particular problem and audience [3]. As interpretability and scientific machine learning enter a new era, researchers must ground themselves in real-world problems and work closely with domain experts.
This paper focuses on scientific machine learning-providiigg dience into a chosen scientific problem- and from its outset, was designed to solve a particularly challenging cosmology problem in close collaboration with cosmologists. We showcase how AWD can inform relevant features in fundamental problem in cosmology: inferring cosmological parame- ters from weak gravitational lensing convergence maps.* 4 In this case, AWD identifies high-intensity peaks in the convergence maps and yields an easily interpretable model which outperforms state-of- the-art neural networks in terms of prediction performance. We next find that AWD successfully provides prediction improvements in another scientific application (now in collaboration with cell- biology experts): molecular-pare partter prediction. In this case, AWD allows to vet that the model's use of clathrin corresponds to our domain knowledge about how clathrin must build up slowly then fall in order to predict successful event. In both cases, the wavelet models from AWD concisely explains model behavior using extremely few parameters (c.g. 10), while also extracting compressed
insight for particular scientific au-
BB attributions, we mean feature importance scores given input data and pre-trained DNN.
*Though we focus on DNNs, AWD works for any black-box models for which we can attain attributions. For the purpose of this work, work with simulated lensing maps.

---

representations of the input in comparison to standard wavelet model. We hope that the depth and grounding of the scientific problems in this work can spur further interpretability research in real- world problems, where interpretability can be evaluated by and enrich domain knowledge, beyond benchmark data contexts such as MNIST [31] where the need for interpretability is less cogent.
2 Background on wavelet transform and TRIM
2.1 Wavelet transform
Wavelets are a class of functions that are localized both in the time and frequency domains. In the classical setting, each wavelet is variation of single wavelet v,, called the mother wavelet. A family of discrete wavelets can be created by scaling and translating the mother wavelet in discrete increments:
{ (Vsm(t)= 3 ('-2") 2
(t-2'n)
(j.)Ez)
where each wavelet in the family vy (1) represents a unique scale and translation of s. With carefully constructed wavelet (see Appendix A.2), the family of wavelets (1) forms orthonormal basis of (IB). Namely, any signal a E (B) can be decomposed into
Σ ΣΣO d im]es,..
(2)
where the wavelet (or detail) coefficients d, [n] at scale 2 are computed by taking the inner product with the basis functions, d; [n] (I, j)...) = x=(t))),nn dt. The decomposition (2) requires an infinite number of scalings to calculate the discrete wavelet transform. To make this decomposition computable, the scaling function á is introduced so that
x Σ as(n]os + ΣΣ dnln
(3)
where (1) = 2--10002--1--- n2) represent different translations of á at scale 2 and [m]m]] (I, JJ,)) are the corresponding approximation coefficients. Conceptually, the ΦJ,n form an orthogonal basis of functions that are smoother at the given scale 2 and therefore can be used to decompose the smooth residuals not captured by the wavelets [32].
fundamental property of the discrete wavelet transform is that the approximation and detail coefficients at scale 2 can be computed from the approximation coefficients of the previous scale 2 [33, 34]. To see this, let us define the two discrete filters, lowpass filter and highpass filter g
1 ( /2 ( (t/12), n)).
Mn| 2 6t//2), nn) and g gin]
Then the following recursive relations hold between the approximation and detail coefficients at two consecutive resolutions:
aa [p] Σ,, Mn 2pjay [n] a > h[pp]: dd [p] = -. En 2pja, [m] a, * g[pp].
(5)
where we denote Mi]]] = M[=n] and g(m] = gو gi- Conversely, the approximation coefficients at scale 2 can be recovered from the coarser-scale approximation and detail coefficients using
[n]as+a[n] 2n Σg
a
Together, these recursive relations lead to the filter bank algorithm, the cascade of discrete convolution and downsampling, which can be efficiently implemented in time O (Signal length). The discrete wavelet transform can be extended two dimensions, using separable (row-column) implementation of ID wavelet transform along each axis (see Appendix A.
3

---

2.2 Transformation Importance (TRIM)
The work here requires the ability to compute attributions which identify important features given input data and trained DNN. Most work on interpreting DNNs has focused on attributing importance to features in the input space of a model, such as pixels in an image or words a document [35-39]. Instead, here we rely on TRIM (Transformation Importance) [40], an approach which attributes importance to features in a transformed domain (here, the wavelet domain) via a straightforward model reparametcrization.
Formally, let f be a pre-trained model that we desire to interpret. If V X - W bijective mapping that maps an input to new domain w = V (z) E W. TRIM reparameterizes the model the case that V is not exactly invertible, TRIM adds the residuals to the output of 1/- ...... is reparameterized by f"(w) S(V-'v w + r) where v = (ax) and r = x (w)) If indexes a subset of features the transformed space indicating which part of the transformed input to interpret, we then define
as f f -- where V denotes the inverse of
TRIM (ws) attr( ws),
where attr(; w) is an attribution method that is evaluated at w and outputs importance value, and where Ws denotes the subvector of w indexed by S. The choice the attribution method attr() can be any local interpretation technique (e.g. LIME [35] or Integrated Gradients (IG) [37]); here focus mainly on the saliency map [41], which simply calculates the gradient of the model's output with respect to its transformed input define feature attribution. We leave more complex attribution methods such as IG or ACD [38] to future work.
3 Adaptive wavelet distillation through interpretations
Adaptive wavelet distillation (AWD) aims to learn a wavelet transform which effectively represents the input data, as well capture information about model trained to predict a response using the input data. Whether or not the resulting wavelet model is (i) sufficiently interpretable and (ii) predictive depends on the context of the problem. Here, we provide two scientific data problems where wavelet models satisfy both criteria (Sec 4).
We now detail how AWD wavelets can be built upon to form an extremely simple model in various contexts (see Sec 4). We first require that the wavelet transform is invertible, allowing for recon- struction of the original data. This ensures that the transform does not lose any information in the input data. We next assure that the learned wavelet is a valid wavelet: the wavelet function 1é and the corresponding scaling function o span sequence of subspaces satisfying the multiresolution axioms [42]. Finally, we add the distillation part of AWD. We calculate the attribution scores of a given model for each coefficient in the wavelet representation, and try find wavelet function zi that makes these attributions sparse. Intuitively, this ensures that the learned wavelet should find representation which can concisely explain a model's prediction. Writing the discrete wavelet transform using the discrete filters h and g (see Eq. 4), we now give final optimization problem for AWD:
minimize C(h.g) k.g
-
Σ c. - m
Ele
2111 m
1
Σ
w Wht (h.g, + Σ [ (8)
Reconstruction less
Wavelet loss
Interpretation less
where Ti denotes wavelet transform operator induced by th and Tt denotes the reconstruction of the data point z4. Here >, > represent hyperparameters that are tuned by users. The only parameters optimized are the lowpass filter h and the highpass filter g. The corresponding scaling and wavelet
functions can be obtained from (h, via the following mapping [32]: Ō(w) = I k(2"Pw) and (ww) - T (2/))((w/2)) where @ and & represent the Fourier transforms of di and va respectively.
Wavelet loss The wavelet loss ensures that the learned filters yield a valid wavelet transform. In contrast to the wavelet constraints used in [11], our formulation introduces additional terms that ensure almost sufficient and necessary conditions on the filters (h, g) to build an orthogonal wavelet basis. Specifically, [32, Theorem 7.2] states the following sufficient conditions on the lowpass filter:

---

if h satisfies
Σ h[n] v2 and [k()]{ + lh( + x)|² 2 for all w,
(9)
well some mild conditions, it can generate scaling function such that the scaled and translated family of the scaling function forms an orthonormal basis of the space of multiresolution approxi- mations of L (R). [43, Theorem 3] further shows that the orthogonality of translates of the scaling function implies that the lowpass filter is orthogonal after translates by 2, i.e.,
{ 0
k = 0
Σ h[m]h[n - 2k]
and as result, || = 1.
(10)
otherwise
Hence the conditions (9), (10) characterize the almost sufficient and necessary conditions on the
lowpass filter. Moreoever, [32, Theorem 7.3] shows that the valid highpass filter can be constructed
from the lowpass filter: in the time domain, it can be written as
g/n| (-1) 11n]], n)..
(11)
where N is the support size of h. Together with (10), we can also deduce that the highpass filter has mean zero, i... Σn g/n] which is necessary for the filter g. See Appendix A.2 for further details.
Finally, we want the learned wavelet to provide sparse representations so we add the l norm penalty on the wavelet coefficients. Combining all these constraints via regularization terms, we define the wavelet loss at the data point I; as
Wh,, xy: Va,l1
h[n] 22)² +
Σ g|ml)² + (1111) 1)²
Σ (((ww) + [lw + +T) 2)² 10
2)² + E h/ml/(n 2k] - 1k=0)²
where g is set as in (11) and X > 0 controls strength of the sparsity of the wavelet representations. We enforce the penalty (h(w))² + lh( + +² 2)², only at the discrete values of w E (2nk k = 1..... N through the discrete Fourier transform. Notice that the wavelet loss does not introduce any additional hyperparameters besides A. In fact, we empirically observe that the sum of penalty terms, except the sparsity penalty, remains very close to zero as long as the filters (h, g) are initialized using known wavelet filters and the interpretation loss is not enforced too strongly.
N
Interpretation loss The interpretation loss enables the distillation of knowledge from the pre- trained model f into the wavelet model. It ensures that attributions in the space of wavelet coefficients Va; are sparse, where the attributions of wavelet coefficients calculated by TRIM [40], as described in Sec 2.2. This forces the wavelet transform to produce representations that concisely explain the model's predictions at different scales and locations. The hyperparameter ? controls the overall contribution of the interpretation loss; large values of y can result in large numerical differences from satisfying the conditions of the mathematical wavelet filters. To our knowledge, this is the first method which uses interpretations from pre-trained model to improve a wavelet representation. This enables the wavelets to not only adapt to the distribution of the inputs, but also gain information about the predicted outputs through the lens of the model f.
4
AWD improves interpretability, prediction performance, and compression in two scientific problems and in simulations
Fig 2 shows visual schematic of the distillation and prediction setup for synthetic and two scientific data problems in this section, whose details will be discussed in the following subsections. $ In both scientific problems, we build extremely simple models based on AWD which significantly outperform the sfate-of-the-art DNN performance with far fewer number of parameters."
'In all experiments, the wavelet function is computed from the corresponding lowpass filter using the Py Waveeets package [44] and building on the Pytorch Wavelets [45, Chapter 3] package. 'For example, the final molecular-parte partner prediction model contains only 10 parameters for the low-pass filter, along with only 30 coefficients in the sparse linear model. The final distilled cosmology model leams only 10 parameters for the low-pat pass filter to make its predictions.
5

---

M of Ma My
Distilled wavelet
Groundtruth wavelet
Synthetic data
Fully connected net
dN
9
d d.
B
LSTM
Distilled wavelet
Linear model on largest coe offciennts
CME
C
Cosmologicul maps
Resnet
Distilled wavelet
Peak counting predicts fmm
Figure Distillation and prediction setup for the three scenarios Sec 4. (A) synthetic simulations, AWD is able to recover groundtruth wavelet (DB5) that are linked to response variable (Sec 4.1). (B) Wavelets distilled by AWD from an LSTM trained to predict molecular partners capture biologically meaningful properties of a large build up in clathrin fluorescence followed by sharp drop and enable prediction using only few key coefficients (Sec 4.2). (C) AWD finds wavelets that are efficient at capturing cosmological information in weak lensing convergence maps and can improve state-of-ie the- art performance cosmological parameter inference using an AWD-based simple peak-counting algorithm (Sec 4.3). Note that the model forms in (B) and (C) come from knowledge about the domain problem.
4.1 Synthetic data
We begin our evaluation using simulations verify whether AWD can recover groundtruth wavelets from noisy data. In these simulations, the inputs are generated i.i.d. from a standard Gaussian distribution AO0 1). To generate the response variable, the inputs domain using Daubechies (DB) 5 wavelets [46], and the response is generated from linear model Y1 v,, 8) + ε,, where the regression coefficients are 2 for few selected locations at particular scale and () otherwise; the noise is generated i.i.d. from a Gaussian distribution (0, 0.1²). Then, 3-layer fully connected neural network with ReLU activations is trained on the pairs of 2s, to accurately predict this response. Note that for any non-singular matrix A, the mapping x (A (A-lur. Vv, groundtruth wavelet explain the model's 's prediction most concisely. The challenge is then to accurately distill the groundtruth wavelet (DB5) from this DNN. This task is fairly difficult: AWD must not only select which scale and locations are important, it must also precisely match the shape of h and g to groundtruth wavelet.
transformed into the wavelet
A (3) predicts the response equally well, but the representations in the
Fig 3 shows the performance of AWD on this task. We initialize the AWD lowpass filter to different known lowpass filters corresponding DBS+random noise (support size 10), Sym- let 5 (support size 10), and Coiflet (support size 12), as shown in Fig 3(A), and then min- imize the objective in Eq. and that minimize the distance to the groundtruth wavelet v/* Distance is measured via d(ú, <l+t) min ming v/ 12 ming || - 112), where zith is the wavelet y circular shifted by and T the wavelet zi flipped in the left/right direction. That is, d calculates the minimum & distance between two wavelets under circular shifts and left/right flip. When the two wavelets have different size of support, shorter wavelet zero-padded to the length of the longer [11]. Fig 3(B) shows that for each different initialization, we find that the distilled wavelet gets considerably closer the groundtruth wavelet. In particular, the results for DBS+noise and Coiflet are nearly identical to the groundtruth and cannot be distinguished in the plot. This is particularly impressive since the
order to recover the groundtruth, we select hyperparameters

---

0.75
DB5+noise
DBS+noise
Sym5
Sym5
0.50
Coif2
Coif2
value
Groundtruth
0.25
0.00
-0.25
10
12
time
Figure 3: AWD accurately identifies the groundtruth important wavelet in simulated data. (A) Plots of the initial lowpass filters. (B) Final wavelets extracted by AWD.
support size of Coiflet differs from that of the groundtruth wavelet, making the task more difficult. Overall, these results demonstrate the ability of AWD to distill key information out of pre-train pre-trained neural network.
4.2 Molecular partner-prediction for a central process in cell biology
We now turn our attention crucial question in cell biology related to the internalization of macromolecules via clathrin-medianee endocytosis (CME) [47]. CME is the primary pathway for entry into the cell, making essential to eukaryotic life [48]. CME an orchestra consisting of hundreds of different protein dynamics, prompting line of research aiming to better understand this process [49]. Crucial to understanding CME is the ability readily distinguish whether or the recruitment of certain molecules will allow for endocytosis, i.e., successfully transporting an object into a cell. Previous approaches have largely relied on the presence marker during imaging (50,51], Alternatively, previous works use domain knowledge to hand- engineer features based on the lifetime an event or thresholds on the recruited amplitude of the clathrin molecule [32,53].
specific scission/uncoating
Here, we aim identify successful CME events with learning approach, obviating the need for an auxiliary marker or hand-engineered features. We use recently published dataset [50] which tags two molecules: clathrin light chain A, which is used as the predictor variable, and auxilin 1, the target variable. In this context, clathrin is used to track the progress of an event, (as recruitment of clathrin molecules usually precedes scission) and recruitment of auxilin molecules follows only when endocytosis successfully occurs (to facilitate disassembly of the clathrin-coated vesicle). See data details in Appendix C. Time-series of fluorescence amplitudes (see Fig 2B) are extracted from raw cell videos for clathrin [52] and used predict the mean amplitude of the auxilin signal, an indicator of whether an event was successful or not. The datasct randomly split into training set of 2,936 data units of dimension 40 and a test set of 1,005 data units. This is challenging problem where deep learning has recently been shown to outperform classical methods. We train a DNN (an L.STM [54]) to predict the auxilin response from the clathrin signal. The model shows state-of-the-art prediction performance, but has extremely poor interpretability and computational cost, so we aim here to distill it into wavelet model through AWD.
Fig shows qualitatively how the learned wavelet function & changes as function of the inter- pretation penalty (increasing to the right) and the sparsity penalty À (increasing downwards). In the initial stage training, we initialize the lowpass filter to correspond to the Daubechics (DB) wavelet. Different combinations of the penalties lead to vastly different learned wavelets, though they all tend to reveal edge-detecting characteristics for a reasonable range of hyperparameter values.
We now test the distilled wavelets for their predictive power. To create an extremely transparent model, we extract only the maximum wavelet coefficients at each scale. By taking the maximum coefficients, these features are expected to be invariant to the specific locations where CME event occurs in the input data. This results in final model with coefficients (6 wavelet coefficients at scales). These wavelet coefficients are used to train a linear model, and the best hyperparameters are selected via cross-validation on the training set. Fig shows the best learned wavelet (for one particular run) extracted by AWD corresponding to the setting of hyperparameters À 0.005
7

---

Increasing interpretation penalty y
00:11
2004
sogm
63031
SAND
3000
Increesinn spaa peraltt
Figure 4: Varying sparsity and interpretation penalty yields different valid wavelets. Wavelet highlighted in red is selected by cross-validation and yields the best prediction performance.
and y 0.043. Table compares the results for AWD to the original LSTM and the initialized, non-adaptive DB5 wavelet model, where the performance is measured via standard R score, proportion of variance in the response that is explained by the model. The AWD model not only closes the gap between the standard wavelet model (DB5) and the neural network, it considerably improves the LSTM's performance (a 10% increase of R score). Table also includes predictive performance of the adaptive wavelet when the interpretation loss is removed during the training; it still outperforms the baseline wavelet (DB5), but fails to outperform LSTM. Moreover, we calculate how well the AWD wavelet compresses its input. Specifically, we measure the proportion of wavelet coefficients in the test set in which the magnitude and the attributions are both above 10-3, The AWD wavelet exhibits much better compression than DB5 (an 18% reduction), showing the ability of AWD to simultaneously provide sparse representations and explain the LSTM's predictions concisely. The AWD model also dramatically decreases the computation time at time, more than 200-fold reudction when compared to LSTM.
In addition improving prediction accuracy, AWD enables domain experts their experimental pipelines by making them more transparent. By inspecting the learned wavelet, AWD allows for checking what clathrin signatures signal successful CME event; it indicates that the distilled wavelet aims to identify large buildup in clathrin fluorescence (corresponding to the building of clathrin-coated pit) followed by sharp drop in clathrin fluorescence (corresponding the rapid deconstruction of the pit). This domain knowledge is extracted from the pre-trained LSTM model by AWD using only the saliency interpretations the wavelet space.
Table 1: Performance comparisons for different models in galecular-partnee prediction. AWD substantially improves predictive accuracy, compression, and computation time on the test set. A higher R2 score, and lower compression factor, and lower computation time indicate better results. For AWD, values are averaged over 5 different random seeds.
AWD
Standard Wavelet (DB5)
LSTM
AWD w/o in- terp. loss
Regression (R² score)
0.262 (0.001)
0.197
0.237
0.231 (0.001) 0.651 (0.003)
Compression factor
0.574 (0.010)
0.704
N/A
Computation time 0.0002s
0.0002s
0.0449s
0.0002s
4.3
Estimating fundamental parameter surrounding the origin of the universe
We now focus on a cosmology problem, where AWD helps replace DNNs with more interpretable alternative. Specifically, we consider weak gravitational lensing convergence maps, i.c.. maps of the mass distribution in the universe integrated up to a certain distance from the observer. In cosmological experiment (e.g. galaxy survey), these mass maps are obtained by measuring the distortion of distant galaxies caused by the deflection of light by the mass between the galaxy and observer [55]. These maps contain wealth physical information matter density in the universe, Pmn Current cosmology research aims to identify the most informative features these maps for inferring the cosmological parameters such as Ωrn The traditional summary statistic for lensing maps is the power spectrum which is known to be sub-optimal for parameter
interest, such as the total

---

inference. Tighter parameter constraints can be obtained by including higher- higher-order statistics, such as the bispectrum [56] and peak counts [57]. However, DNN-based inference methods claim to improve on constraints based on these traditional summaries [58-60].
Here, we aim improve the predictive power of DNN-based methods while gaining interpretability by distilling a predictive AWD model. In this context, it is critically important to obtain interpretability. as it provides deeper understanding into what information most important to infer Shn, and can be used to design optimal experiments or analysis methods. Moreover, because these models are trained on numerical simulations (realizations of the Universe with different cosmological parameters), is important to validate that the model uses reasonable features rather than latching on to numerical artifacts the simulations. We start by training model to predict Q' from simulated weak gravitational lensing convergence maps. We train a DNN to predict S from 100,000 mass maps simulated with 10 different sets cosmological parameter values at the universe origin from the Massi VeNuS simulations [62] (full simulation details given in Appendix D), achieving an R² value of 0.92 on the test set (10,000 mass maps); Fig 2C shows an example mass map.
We again construct an interpretable model using the wavelets distilled by AWD from the trained DNN, To make predictions, we use the simple peak-counting algorithm developed in a previous work [59], which convolves a peak-finding filter with the input images. Then, these peaks are used to regress on the outcome. In contrast to the fixed filters such as Laplace or Roberts cross used in previous works [59], here we use the wavelets distilled by AWD, which result in three 2D wavelet filters (LL, LH, HL) and the 2D approximation filter (LL). The size of the distilled AWD filters 12x12 and inspection of these filters shows majority of the mass is concentrated on 3x3 subfilters (see Fig 2C). Then we extract those subfilters to use for
filters- by doing so, the size of the filters match with those used in [59] (additional details given in Appendix D. 1). The hyperparameters for AWD are selected by evaluating the predictive model's performance on a held-out validation set.
Table 2 shows the results of predicting using the peak-finding algorithm with various filters. The evaluation metric is the RMSE (Root mean square error). Its performance again outperforms the fully trained neural network (Resnet) model and the standard non-adaptive wavelet (DB5) model, as well as other bascline methods using Laplace filter and Roberts cross filter (see Appendix D.I for details on how these filters are defined). Without the interpretation loss, the adaptive wavelet fails to outperform the Roberts cross filter and Resnet model. Moreover, as can be seen in the compression factor, the AWD wavelet provides more efficient representations for the mass maps as well as concise explanation for the DNN's predictions compared to the DB 5 wavelet.
Table 2: Performance comparisons for different models in cosmological parameter prediction. The lower RMSE and compression factor indicate better results. For RMSE, standard deviations are estimated from 10, 000 bootstrap samples.
AWD
Roberts- Cross
Laplace
DB5 Wavelet
Resnet
AWD
w/o
interp. loss
Regression (RMSE x00=
1.029
1.259
1,369
1.569
1.156
1.354
(0.033)
(0.039)
(0.047)
(0.048)
(0.024)
(0.047)
Compression factor
0.610
N/A
N/A
0.620
N/A
0.616
Fig 2C shows the learned AWD filters corresponding to the best distilled wavelet. The learned wavelet filters are symmetric and resemble the "matched filters" which have been used in the past to identify peaks on convergence maps in the cosmology literature [63, 64]. We expect from cosmology knowledge that much information is contained in the peaks of the convergence maps (their amplitude, shape, and numbers), so this indeed matches our expectations based on physics. The high predictive performance further demonstrates that the AWD filters are more efficient at capturing cosmological information and better adapted to the shape of the peaks, than standard wavelets could do.
Moreover, the adaptive wavelet distillation allows us look at "wavelet activation maps" (see Fig D2) to localize where in the convergence maps important information is concentrated. In other words,
The model's architecture is Resnet 18 [61], modified take only one input channel.

---

we can indeed see that the AWD wavelet concentrates on identifying high intensity peaks, which is where most of the "localized" information expected from theory.
Discussion
In this work, we introduce AWD, a method to distill adapative wavelets from a pre-trained supervised model such as DNNs for interpretation. Doing so enables AWD to automatically detect and adapt to aspects of data that are important for prediction in an interpretable manner. The benefits of distilling relevant predictive information captured in a DNN are demonstrated through applications to synthetic and real data in two scientific settings. Overall, AWD allows us to interpret a DNN in terms of conventional wavelets, bringing interpretability with domain insights while simultaneously improving compression and computational costs, all while preserving or improving predictive power.
Limitations and future work AWD works well in particular domains where wavelets are rea- sonable modeling choice, such as images and time-series, which possess multi-scale structure. Other domains, such DNA-sequence data or text data which do not posses this structure would not benefit from AWD. Therefore, it is important to use domain knowledge to decide whether AWD can be beneficial for a given task and domain. We also used domain knowledge pick reasonable choices for the number of scales in the distilled wavelet transform, and achieved strong performance without tuning this parameter. Future work could include more analysis on the effect of the number of scales, or an optimization procedure to pick it.
In this paper we test method with the saliency attribution method; however, many alternative interpretation techniques have been proposed recently, such as SHAP [36] or Contextual Decom- position [38, 65], and the comparison between different interpretation techniques can be carefully explored in the context of particular problem and audience. Care must be taken to ensure that these methods can be efficiently optimized, as even the simplest saliency methods requires taking a second partial derivative during distillation, which is computationally expensive especially for large data and network sizes.
Throughout the experiments, the learned wavelets are initialized to a filter corresponding to a known wavelet. Initializing to a known wavelet filter makes the optimization faster and more stable, but when initializing randomly, the predictive performance can be unstable, especially when the variance be solved by initializing at any known wavelet
of the initialization is high. However, this (potentially with some random noise added to it) or first performing few steps of the optimization without the interpretation loss before adding it in. Another promising strategy is initializing with an (unsupervised) dictionary of wavelets [19].
The current work learns single-layer wavelet transform, but the complex nature of modern datasets often require strong nonlincaritics. Future work could extend AWD beyond a single-layer wavelet transform, c.g. by borrowing ideas from scattering transform [27] or to other interpretable models [2, 66]. This would allow for bridging closer to deep learning while keeping interpretability, which can be effectively applied to other areas, such computer vision and natural-image classification. We hope to continue this line of research in order to improve the interpretability and computational efficiency of DNN models across many domains ranging from physical and biomedical sciences to computer vision and information technology.
Societal impacts AWD generally helps alleviate many concerns around DNNs, both in terms of interpretability and computational resources. However, its optimization at training time necessitates training a DNN, which can incur high computational cost (and therefore corresponding impacts from energy usage). Additionally, the interpretability granted by AWD could lead to over-reliance on the fitted models, even in applications where it does not perform well.
Acknowledgements
We gratefully acknowledge partial support from NSF TRIPODS Grant 1740855, DMS- 1613002, 1953191, 2015341, IIS 1741340, ONR grant 17-014-17-1-2176. Moreover, this work is supported in part by the Center for Science of Information (CSol), an NSF Science and Technology Center, under grant agreement CCF-0939370, by the NSF Grant DMS 203 1883 "Collaboration the Theoretical
10

---

Foundations of Deep Learning", and by the NSF Foundations of Data Science Institute (FODSI), SU was supported with funding from Philomathia Foundation and Chan Zuckerberg Initiative Imaging Scientist program. The authors would also like to thank Alan Dong for enlightening discussions on optimizing wavelets. The authors would also like to thank Tom Kirchhausen, Kangmin He, Eli Song, and Song Dang for providing the clathrin mediated endocytosis data to apply AWD for molecular partner predictions. We would like to gratefully acknowledge AWS computing credits as part NSF Award 1741340.
References
[1] Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature Machine Intelligence, (()):26-2155, 2019. (Cited on page 1.)
[2] Cynthia Rudin, Chaofan Chen, Zhi Chen, Haiyang Huang, Lesia Semenova, and Chudi Zhong. Interpretable machine learning: Fundamental principles and 10 grand challenges. arXiv preprint 1rXiv:213.112511, 2021. (Cited on pages and 10.)
[3] W James Murdoch, Chandan Singh, Karl Kumbier, Reza Abbasi-Asl, and Bin Yu. Definitions, methods, and applications in interpretable machine learning. Proceedings of the National Academy of Sciences, 116(44):22071-22 1-22080, 2019. (Cited on pages and 2.)
[4] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprini arXiv: /503.02531, 2015. (Cited page 2.)
[5] Guodong Xu, Ziwei Liu, and Chen Change Loy. Computation-efficiett knowledge distillation via uncertainty-aware mixup. arXiv preprint arXiv:2012.09411 2020. (Cited on page 2.)
[6] Jia Guo, Minghao Chen, Yao Hu, Chen Zhu, Xiaofei Hc, and Deng Cai. Spherical knowledge distillation. arXiv preprint arXiv:2010.07485, 2020. (Cited page 2.)
[7] Nicholas Frosst and Geoffrey Hinton. Distilling neural network into soft decision tree. arXiv preprint Xrviv: /7 I. .07784,7 2017. (Cited on page 2.)
[8] Mark Craven and Jude W Shavlik. Extracting tree-structured representations of trained networks. In Advances in neural information processing systems, pages 24-30, 1996. (Cited on page 2.)
[9] Jiawei Li, Yiming Li, Xingchun Xiang, Shu-Tao Xia, Siyi Dong, and Yun Cai. Tnt: An interpretable tree-network-tee learning framework using knowledge distillation. Entropy, 22(1 1203, 2020. (Cited on page 2.)
[10] Sarah Tan, Rich Caruana, Giles Hooker, Paul Koch, and Albert Gordo. Learning global additive explanations for neural nets using model distillation. arXiv preprint 110iv:1801.080 2018. (Cited on page 2.)
[11] Daniel Recoskie and Richard Mann. Learning sparse wavelet representations. arXiv preprint (:Xiv:180200296: 2018. (Cited pages 2, 4, and 6.)
[12] Phil Sallee and Bruno Olshausen. Learning sparse multiscale image representations, Advances in neural information processing systems, 15:1351-1358, 2002. (Cited on page 2.)
[13] Jing Lin and MJ Zuo. Gearbox fault diagnosis using adaptive wavelet filter. Mechanical systems and signal processing, 1766):1259-1269 2003. (Cited on page 2.)
[14] Laurent Brechet, Marie-Françoiss Lucas, Christian Doncarli, and Dario Farina. Compression of biomedical signals with mother wavelet optimization and best-basis wavelet packet selection. IEEE Transactions on Biomedical Engineering, 54(12):2186-21)
2007. (Cited on page 2.)
[15] James M Hereford, David W Roach, and Ryan Pigford. Image compression using parameterized wavelets with feedback. In Independent Component Analyses, Wavelets, and Neural Networks, volume 5102, pages 267-277. International Society for Optics and Photonics, 2003. (Cited on page 2.)
[16] Daniel Recoskie. Learning sparse orthogonal wavelet filters. 2018. (Cited page 2.)
[17] Dhruv Jawali, Abhishek Kumar, and Chandra Sekhar Seelamantula. A learning approach for wavelet design. In ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 5018-5022 IEEE, 2019. (Cited on page 2.)

---

[18] Cheng Tai and E Weinan. Multiscale adaptive representation of signals: I. the basic framework. The Journal of Machine Learning Research, 17(1):4875-49) 2016. (Cited on page 2.)
[19] Boaz Ophir, Michael Lustig, and Michael Elad. Multi-scale dictionary learning using wavelets. IEEE Journal of Selected Topics in Signal Processing, 5(5):1014-1024, 2011. (Cited on pages 2
and 10.)
[20] Qinghua Zhang and Albert Benveniste. Wavelet networks. IEEE transactions on Neural Networks, 3(6):889 898, 1992. (Cited on page 2.)
[21] Jingyuan Wang, Ze Wang, Jianfeng Li, and Junjie Wu. Multilevel wavelet decomposition network for interpretable time series analysis. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2437-2446, 2018 8.
(Cited on page 2.)
[22] Haidar Khan and Bulent Yener. Learning filter widths of spectral decompositions with wavelets. Advances in Neural Information Processing Systems, 31:4601 116600146122 12, 2018. (Cited on page 2.)
[23] Jun Zhang, Gilbert Walter, Yubo Miao, and Wan Ngai Wayne Lee. Wavelet neural networks for function learning. IEEE transactions on Signal Processing, 43(6):1455-1497) 1995. (Cited
on page 2.)
[24] J. Dheeba, N. Albert Singh, and S. Tamil Selvi. Computer-aided detection of breast cancer on mammograms: A swarm intelligence optimized wavelet neural network approach. Journal of Biomedical Informatics, 49:45-52, 2014. (Cited on page 2.)
[25] Tianfu Li, Zhibin Zhao, Chuang Sun, Li Cheng, Xuefeng Chen, Ruqiang Yan, and Robert X Gao. Waveletkernelnet: An interpretable deep neural network for industrial intelligent diagnosis. IEEE Transactions on Systems, Man, and Cybernetics: Systems, 2021. (Cited on page 2.)
[26] Stéphane Mallat. Group invariant scattering. Communications on Pure and Applied Mathematics, 65(10):1331-1399 2012. (Cited on page 2.)
[27] Joan Bruna and Stéphane Mallat. Invariant scattering convolution networks. IEEE transactions on pattern analysis and machine intelligence, 35(8):1872-1886, 2013, (Cited on pages and 10.)
[28] Christoph Molnar. Interpretable machine learning. Lulu. com, 2019. (Cited on page 2.)
[29] Julius Adebayo, Justin Gilmer, Michael Muclly, Ian Goodfellow, Moritz Hardt, and Been Kim. Sanity checks for saliency maps. In Advances in Neural Information Processing Systems, pages 9505-9515, 2018. (Cited on page 2.)
[30] Finale Doshi Velez and Been Kim. roadmap for a rigorous science of interpretability. arXiv preprint arXiv: / 08608, 2017. (Cited on page 2.)
[31] Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mnist//
1998. (Cited on page 3.)
[32] Stéphane Mallat. A wavelet tour of signal processing, Third edition: The sparse way. Academic Press, 2008. (Cited pages 3, 4,5, 15, and 16.)
[33] Stephane Mallat. Multiresolution approximations and wavelet orthonormal bases of (R). Transactions of the American mathematical sociery, 315(1):69-88) (:69-77, 1989. (Cited on pages
and 15.)
[34] Yves Meyer. Wavelets and Operators: Volume 1. Number 37. Cambridge university press, 1992. (Cited on pages 3 and 15.)
[35] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. Why should trust you?: Explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 1135-1144. ACM, 2016. (Cited
on page 4.)
[36] Scott M Lundberg and Su-In Lee. A unified approach to interpreting model predictions. In Advances in Neural Information Processing Systems, pages 4777, 2017. (Cited on pages 4
and 10.)
[37] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. ICML, 2017. (Cited on pages and 22.)
[38] Chandan Singh, W. James Murdoch, and Bin Yu. Hierarchical interpretations for neural network predictions. In International Conference on Learning Representations, 2019. (Cited on pages
and 10.)
12

---

[39] Laura Rieger, Chandan Singh, W James Murdoch, and Bin Yu. Interpretations are use- ful: penalizing explanations align neural networks with prior knowledge. arXiv preprint arXiv: /909. 13584, 2019. (Cited on page 4.)
[40]
Chandan Singh, Wooseok Ha. Francois Lanusse, Vanessa Bochm, Jia Liu, and Bin Yu. Trans- formation importance with applications to cosmology. arXiv preprint arXiv:2003.01926 2020. (Cited on pages 4 and 5.)
[41] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv:/312.00 131.60334 2013. (Cited on page 4.)
[42] Stephane Mallat. A theory for multiresolution signal decomposition: the wavelet representation. IEEE transactions pattern analysis and machine intelligence, 11(7):674-69) 693, 1989. (Cited on page 4.)
[43] Sidney Burrus, Ramesh Gopinath, and Haitao Guo. Introduction to wavelets and wavelet transforms: a primer. Englewood Cliffs, 1997. (Cited on pages and 16.)
[44] Gregory Lee, Ralf Gommers, Filip Waselewski, Kai Wohlfahrt, and Aaron O'Leary. Pywavelets: A python package for wavelet analysis. Journal of Open Source Software, 4(36): 2019. (Cited on page 5.)
[45] Fergal Cotter. University of Cambridge, 2020. (Cited on page 5.)
Uses of Complex Wavelets in Deep Convolutional Neural Networks. PhD thesis,
[46] DAUBECHIES, Orthonormal bases of compactly supported wavelets. Commun. Pure Appl. Math. 41:909-996, 1988. (Cited on page 6.)
[47] Tom Kirchhausen, David Owen, and Stephen C Harrison. Molecular structure, function, and dynamics of fathrin-mediated membrane traffic. Cold Spring Harbor perspectives in biology, 6())a016722, 2014. (Cited on page 7.)
[48] Harvey T McMahon and Emmanuel Boucrot. Molecular mechanism and physiological functions of (Cited on page 7.)
endocytosis. Nature reviews Molecular cell biology, 12(8):517, 2011.
[49] Marko Kaksonen and Aurélien Roux. Mechanisms of clathrin endocytosis. Nature Reviews Molecular Cell Biology, 19(5):313, 2018. (Cited on page 7.)
[50] Kangmin He, Eli Song, Srigokul Upadhyayula. Song Dang. Raphael Gaudin, Wesley Skillem, Kevin Bu, Benjamin R Capraro, Iris Rapoport, Ilja Kusters, et al. Dynamics of auxilin 1 and gak in Nathrin-mediated traffic. Journal of Cell Biology, 219(3), 2020. (Cited on pages and 18.)
[51] Xinxin Wang, Zhiming Chen, Marcel Mettlen, Jungsik Noh, Sandra L Schmid, and Gaudenz Danuser. Dasc, sensitive classifier for measuring discrete early stages in clathrin-mediated endocytosis. eLife, 9:e53686, 2020. (Cited on page 7.)
[52] François Aguet, Costin N Antonescu, Marcel Mettlen, Sandra L Schmid, and Gaudenz Danuser. Advances in analysis of low signal-to-noise images link dynamin and ap2 to the functions of an endocytic checkpoint. Developmental cell, 26(3):279-295 29177999911199 2013. (Cited on pages 7 and 18.)
[53] Zuzana Kadlecova, Stephanic J Spielman, Dinah Loerke, Aparn Mohanakrishnan, Dana Kim Reed, and Sandra L Schmid. Regulation of Blathrin-mddiane endocytosis by hicrarchical allosteric activation of ap2. Journal of Cell Biology, 216(1):167-179) 2017. (Cited page 7.)
[54] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation. 183::1735-7780, 1997. (Cited on page 7.)
[55] Matthias Bartelmann and Peter Schneider. Weak gravitational lensing. Physics Reports, 340(4- 9):291-472, 2001. (Cited on page 8.)
[56] William R Coulton, Jia Liu, Mathew S Madhavacheril, Vanessa Böhm, and David N Spergel. Constraining neutrino mass with the tomographic weak lensing bispectrum. Journal of Cosmol- ogy and Astroparticle Physics, 2019(05):043, 2019. (Cited on page 9.)
[57] Zack Li, Jia Liu, José Manuel Zorrilla Matilla, and William R Coulton. Constraining neutrino mass with tomographic weak lensing peak counts. Physical Review D, 9966)22335 2019. (Cited page 9.)
13

---

[58] Dezsö Ribli, Bälint Armin Pataki. José Manuel Zorrilla Matilla, Daniel Hsu, Zoltán Haiman, and István Csabai. Weak lensing cosmology with convolutional neural networks on noisy data. Monthly Notices of the Royal Astronomical Society, 490(2):1843-18) 186-18800, 2019. (Cited on page 9.)
[59]
Dezsö Ribli, Bälint Armin Pataki, and István Csabai. An improved cosmological parameter inference scheme motivated by deep learning. Nature Astronomy, 3(1):93, 2019. (Cited on pages 9, 21, and 22.)
[60] Janis Fluri, Tomasz Kacprzak, Aurelien Lucchi, Alexandre Refregier, Adam Amara, Thomas Hofmann, and Aurel Schneider. Cosmological constraints with deep learning from kids-450 weak lensing maps. Physical Review D, 100(6)006314, 2019. (Cited page 9.)
[61]
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016. (Cited on page 9.)
[62] Jia Liu, Simeon Bird, José Manuel Zorrilla Matilla, J. Colin Hill, Zoltán Haiman, Mathew Madhavacheril, Andrea Petri, and David N. Spergel. MassiveNuS: cosmological massive neutrino simulations. JCAP, 2018(3):049) 2018. (Cited on pages 9 and 20.)
[63] Matteo Maturi, Massimo Meneghetti, Matthias Bartelmann, Klaus Dolag, and Lauro Moscardini. An optimal filter for the detection of galaxy clusters through weak lensing, Astronomy & Astrophysics, 442(3):851 8600, 2005. (Cited on page 9.)
[64]
Fabian Schmidt and Eduardo Rozo. Weak-lensing peak finding: Estimators, filters, and biases.
The Astrophysical Journal, 735(2):119, 2011. (Cited page 9.)
[65] W James Murdoch, Peter Liu, and Bin Yu. Beyond word importance: Contextual decom- position to extract interactions from Istms. arXiv preprint arXiv:1801.05453. 2018. (Cited on page 10.)
[66] Chandan Singh, Keyan Nasseri, Yan Shuo Tan, Tiffany Tang, and Bin Yu. imodels: a python package for fitting interpretable models. Journal of Open Source Software, 6661):3192, 2021. (Cited on page 10.)
14


Reviews:
[{"review_id": "fGKSSnnXgYa", "reviewer": "Reviewer_HEh3", "summary": "This paper proposes learning an informative wavelet decomposition of the input space of a neural network that is able to come up with sparse attributions of the features used in the output. The idea is that in some scientific domains which collect time or space oriented data (in the form of images or signals), the features that lead to a particular prediction are hard to interpret because of complex patterns, and that wavelets can instead provide a concise description since they are localized in both frequency and space. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I find this to be a valuable contribution in overcoming the black boxness of neural networks used in scientific domains. The example case studies in biology and cosmology are convincing and clearly written even to those with a limited background in this field.  The summary of the wavelet transform (S2.1) is clear enough to make this paper accessible to a reader who is unfamiliar with wavelets.  However, the overall way the paper is written could be improved.  First, the authors should make it clear what sort of domains benefit from these---clearly structured domains. For example frequency domain features would likely be of little use in interpreting DNA sequences or text data. It also probably more useful in multi-granular settings where despite the high dimensionality of the underlying space, the phenomenon is a signal over the space. From this standpoint, this work could probably also be extended to graph signals. \n\nSecond, the authors don't use much bandwidth to discuss selection of wavelets, scales, or levels (in cases of multilevel transforms), but these are precisely what would be domain-specific, as the authors note that interpretations often are. Indeed it seems as if first learning a useful dictionary of wavelets may be helpful. See Aharon et al. 2006, or Ophir et al. 2011 for wavelet dictionary learning algorithms. \n\nFigure 1 does not seem very helpful. The authors instead should showcase that the transform is trained with several losses including an autoencoder-like loss reconstruction loss, and a sparsity loss which the authors call 'interpretability loss.' This could be visualized along side a neural network. \n\nIn Line 43, the statement that wavelets have not been used for interpretability is not quite true. I believe there are quite a few papers applying wavelet and using it in terms of model interpretability. Just a few sample papers:  \nWang et al., Multilevel Wavelet Decomposition Network for Interpretable Time Series Analysis\nLi et al, WaveletKernelNet: An Interpretable Deep Neural Network for Industrial Intelligent Diagnosis\nKhan et al, Learning filter widths of spectral decompositions with wavelets\n\nAdditional points of clarification: \n\n1) Do the wavelets have to be initialized to a known filter or it can be random? Is the distance loss added during training to recover the ground truth wavelet? \n\n2) In the experimental section, the authors showed the AWD method helped enhance the prediction accuracy, but how is the interpretability compared to other methods? Can the authors provide some comparisons of AWD and other methods so the idea of enhancing interpretability is further illustrated?\n\nI would also suggest the authors to conduct an ablation study to see the effects of each loss term on learning the wavelet transforms? Since the key of this method is the combination of three losses, it’s good to show how these loss terms affect the model.\n\nI am satisfied by the author responses and will maintain my high score. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors should clearly set out the use case for this contribution. Clearly this is only useful for signals in time or space domains. ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 "}, "extra_text": {}}, {"review_id": "S3GKzQZw4-2", "reviewer": "Reviewer_eZkr", "summary": "The paper proposes distilling a DNN into a learnable wavelet transform. Producing an interpretable architecture, that was used in two scientific applications cosmology and molecular-partner prediction.\nThe paper proposed 3 losses: In AWD the wavelet transform should be invertible, allowing for reconstruction of the original data. This ensures that the transform does not lose any information in the input data \"reconstruction loss\".  The distillation part of AWD is added by calculating the attribution scores of a  given model for each coefficient in the wavelet representation and try to find a wavelet function that makes these attributions sparse \"interpretation loss\". Finally, the wavelet loss which ensures that the learned filters yield a valid wavelet transform.\n\nThe paper provides two scientific data problems where the resulting wavelet model is sufficiently interpretable for use and provides similar prediction performance as the DNN.  The paper also evaluated AWD on Synthetic data showing that AWD can recover groundtruth wavelets from noisy data.\n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n- The work is novel, and original.\n- The paper is well written and clear.\n- Significance of work is high since learned wavelets can be used as interpretations for scientific domains.\n- The submission is technically sound.\n\nWeakness:\n- No comparison was done with other distillation methods.\n- Transformation Importance relies on gradients for importance identifications which are noisy.\n\nQuestions:\n- Do you have any insights on why distillation improves accuracy above DNN? Typically the upper bound is DNN accuracy.\n\nTypos:\n- line 106 bijective-> objective\n- line 176 think you mean Fig 2?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "Ld6HmYLrGXl", "reviewer": "Reviewer_XipH", "summary": "The authors propose an extended approach for learnable wavelet transforms, and add a distillation-based regularization term to the loss, which penalizes parts of the wavelet with low importance scores in the wavelet-transformed space given a pre-trained neural network.\nIn experiments on synthetic data, as well as two scientific datasets, their predictive performance is quantitatively analyzed.\nAdditionally, the interpretability of the learned wavelets is analyzed qualitatively.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Originality:\nPrevious approaches to learned sparse wavelet transforms are extended by additional constraints from previous literature.\nThe main originality is the proposed neural network-attribution-based distillation regularization term introduced into the learned wavelet loss.\nThe paper is fairly original, and differentiates itself well from previous work.\n\n### Quality:\nIt is not clear, how much the teacher model actually contributes to the increased performance with the intepretation loss, especially in Section 4.3, where the AWD outperforms its teacher.\nReporting the RSME for the case where the interpretation penalty is set to zero in Tables 1 and 2 may give insight to whether and how much the interpretation contributes to increased performance.\n\nThe authors claim high interpretability of the learned wavelets.\nThis may be highly dependent on the problem domain.\nIntroducing an experiment for which wavelets may not be the best choice, may better show the limitations of the interpretability of method, and make it more convincing.\nWhile it is intuitive, that these learned wavelets give a much clearer idea of their produced predictions, the analysis to support the claim is purely qualitative, and may benefit from a more quantitative analysis.\nSuch an analysis would furthermore improve the quality for the paper, since the two real-world experiments rely on expert knowledge for their qualitative interpretation.\n\n\n### Clarity:\nThe paper is clearly written, the introduced methods are introduced well and sufficiently.\n\nFigure 4 shows qualitatively how different hyperparameters yield different wavelets.\nBeyond the fact that they change, it is not really clear to me how the different wavelets reflect the prediction strategy.\nA discussion of the possibly increased intuitive \"goodness\" of the wavelets for different parameters may increase the clarity of this experiment\n\nMinor:\n- l.121-122 may or not -> may or may not\n\n### Significance:\nAlthough a broader application of the method may increase its significance, the results suggest a valid extension for learnable wavelets.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are not discussed.\nWhile wavelets provide a powerful framework, the authors should discuss in which cases the training may fail.\nThe problem of the generally loose term of \"explanation\" as discussed in the introduction.\nNo negative societal impact is reported.\nMaybe at least a potential societal impact more general terms of interpretability may be addressed.\nFor example, over-confidence in a model caused by over-interpretation of possible red herrings in the wavelets could have a negative impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "-uIGAR4fVF9", "reviewer": "Reviewer_KLEM", "summary": "The authors proposed Adaptive Wavelet Distillation (AWD) for a learnable wavelet transform. It distills the discriminative power of a deep neural network (DNN) by penalizing its feature attributions in the wavelet domain. Specifically, the authors proposed a new loss term consisting of reconstruction loss for invertibility, wavelet loss for yielding a valid transform, and interpretation loss for distilling the information of a trained DNN. Through synthetic, cosmology, and biological datasets, they demonstrated that the proposed method provides SoTA predictive performance and identifies interpretable features that are scientifically meaningful in the corresponding domains. ", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "(1) Originality\n- The authors proposed to learn wavelet transforms based on the attributions from a trained DNN. To the best of my knowledge, this is a novel approach to use the wavelet transforms and knowledge distillation to interpret a trained DNN model.\n\n(2) Quality\n- The authors argue that wavelet models may or may not have better interpretability depending on a problem's context (line 36-38, line 121-123). Can you provide examples where wavelet models have better interpretability? Furthermore, I think explaining why having an orthogonal basis, and sparse representations indicate better interpretability would help understand the benefit of this paper.\n- I am not sure the compression rate (line 243-245, 297-299) can show the better interpretability of the proposed method. In addition, while the authors stated that they had close collaboration with domain experts (lines 10-12), I could not find relevant details in the paper. Simply stating that the learned wavelet resembles a canonical pattern does not provide solid grounds.\n- The interpretation loss induces the wavelet transform to yield the wavelet whose TRIM attribution is sparse concerning a trained DNN model. Since it largely deviates from the standard knowledge distillation setting where prediction outputs of a DNN are utilized, please provide a more detailed rationale for how the interpretation loss enables \"gain information about the predicted outputs.\"\n- It is unclear how critical the interpretation loss is for the proposed learnable wavelet transform. The authors only showed how the learned wavelet function changes in terms of the interpretation penalty (Figure 4). Please provide ablation studies showing how interpretation loss affects the predictive power of AWD. I think AWD without the interpretation loss may provide similar predictive performances and compression rates, which would indicate the distillation of a trained DNN is unnecessary.\n- How critical are the initialization using known wavelet filters for the predictive performance? Is it more essential than the interpretation loss? Please compare with the results of using random initialization. Furthermore, please provide more detailed explanations for the different initialization methods (support sizes, Coiflet2, DB5).\n- The authors claim that AWD outperforms the SoTA neural networks. Are the compared LSTM, and ResNet models are SoTA models? If not, please provide and adequately compare with the SoTA methods.\n- The authors claim that AWD allows us to look at wavelet activation maps (Figure D2, line 307). However, activation maps can also be generated using the feature attribution in the input domain. Please show that whether AWD activation maps provide \"better\" interpretation. In addition, while the manuscript used a saliency map for attribution scores, Figure D2 uses integrated gradients. The inconsistency may lead to misunderstanding.\n\n(3) Clarity\nIn general, the manuscript is clearly written. Another round of proofreading would be helpful, considering that there are some typos. (Ex. Line 79 (L_2), Line 308 Fig. C2). \n\n(4) Significance\nAlthough the paper has its own merits, unfortunately, it seems that both additional quantitative and qualitative analyses are required to claim better interpretability. While I agree that interpretability can be better defined in a particular problem, the manuscript does not seem to define and show the interpretability in the contexts properly. I think it does not fold great significance in the current form.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The main limitation of the paper is that the claimed interpretability is not clearly demonstrated through the experiments. Please refer to the Main Review for more details. ", "review_text": "", "extra_scores": {"ethical_concerns": "I do not see any ethical issues with this paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agreed that this paper had many strengths as a novel method. Presenting two applications as case studies was a welcome departure from standard NeurIPS submissions. The reviewers have given you a few actionable suggestions which I hope you will take for the final version. The idea of including examples where wavelets are not the best choice in the supplementals is a very good one.
Decision: Accept (Poster)