Adaptive wavelet distillation from neural networks through interpretations

1 Introduction . Recent advancements in deep learning have led to 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 highstakes applications such as medicine , biology , and policy-making . In these cases , interpretations which are relevant to a 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 to be faster and more computationally efficient than large neural networks . 1 github.com/Yu-Group/adaptive-wavelets 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . One promising approach to constructing interpretable models without sacrificing prediction performance 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 a global additive model [ 10 ] , with some success . Here , we consider distilling a DNN into a learnable wavelet transform , which is a powerful tool to describe signals both in time ( spatial ) and frequency domains that has found numerous successful applications in physical and biomedical sciences . Wavelets have many properties amenable to interpretation : they can form an orthogonal basis , identify a sparse representation of a 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 1 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 wavelets ; 2 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.3 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 , a major part of the issue is that interpretability can not 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—providing insight for a particular scientific audience 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 a fundamental problem in cosmology : inferring cosmological parameters 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-ofthe-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 cellbiology experts ) : molecular-partner prediction . In this case , AWD allows us 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 a successful event . In both cases , the wavelet models from AWD concisely explains model behavior using extremely few parameters ( e.g . 10 ) , while also extracting compressed 2By attributions , we mean feature importance scores given input data and a pre-trained DNN . 3Though we focus on DNNs , AWD works for any black-box models for which we can attain attributions . 4For the purpose of this work , we work with simulated lensing maps . representations of the input in comparison to a standard wavelet model . We hope that the depth and grounding of the scientific problems in this work can spur further interpretability research in realworld 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 a variation of a single wavelet , called the mother wavelet . A family of discrete wavelets can be created by scaling and translating the mother wavelet in discrete increments : ⇢ j , n ( t ) = 1 p 2j ✓ t 2jn 2j ◆ ( j , n ) 2Z2 , ( 1 ) where each wavelet in the family j , n ( t ) represents a unique scale and translation of . With a carefully constructed wavelet ( see Appendix A.2 ) , the family of wavelets ( 1 ) forms an orthonormal basis of L2 ( R ) . Namely , any signal x 2 L2 ( R ) can be decomposed into x = X n X j dj [ n ] j , n , ( 2 ) where the wavelet ( or detail ) coefficients dj [ n ] at scale 2j are computed by taking the inner product with the basis functions , dj [ n ] = hx , j , ni = R x ( t ) j , n ( t ) 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 = X n aJ [ n ] J , n + X n JX j dj [ n ] j , n , ( 3 ) where J , n ( t ) = 2 J/2 ( 2 J t n ) represent different translations of at scale 2J and aJ [ n ] = hx , J , ni are the corresponding approximation coefficients . Conceptually , the J , n form an orthogonal basis of functions that are smoother at the given scale 2J , and therefore can be used to decompose the smooth residuals not captured by the wavelets [ 32 ] . A fundamental property of the discrete wavelet transform is that the approximation and detail coefficients at scale 2j+1 can be computed from the approximation coefficients of the previous scale at 2j [ 33 , 34 ] . To see this , let us define the two discrete filters , lowpass filter h and highpass filter g h [ n ] = h 1 p 2 ( t/2 ) , ( t n ) i and g [ n ] = h 1 p 2 ( t/2 ) , ( t n ) i . ( 4 ) Then the following recursive relations hold between the approximation and detail coefficients at two consecutive resolutions : ⇢ aj+1 [ p ] = P n h [ n 2p ] aj [ n ] = aj ? h̄ [ 2p ] ; dj+1 [ p ] = P n g [ n 2p ] aj [ n ] = aj ? ḡ [ 2p ] , ( 5 ) where we denote h̄ [ n ] = h [ n ] and ḡ [ n ] = g [ n ] . Conversely , the approximation coefficients at scale 2j can be recovered from the coarser-scale approximation and detail coefficients using aj [ p ] = X n h [ p 2n ] aj+1 [ n ] + X n g [ p 2n ] dj+1 [ n ] . ( 6 ) 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 to two dimensions , using a separable ( row-column ) implementation of 1D wavelet transform along each axis ( see Appendix A.1 ) . 2.2 Transformation Importance ( TRIM ) . The work here requires the ability to compute attributions which identify important features given input data and a 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 in 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 reparameterization . Formally , let f be a pre-trained model that we desire to interpret . If : X ! W is a bijective mapping that maps an input x to a new domain w = ( x ) 2 W , TRIM reparameterizes the model as f 0 = f 1 , where 1 denotes the inverse of . In the case that is not exactly invertible , TRIM adds the residuals to the output of 1 , i.e. , f 0 is reparameterized by f 0 ( w ) = f ( 1w + r ) where w = ( x ) and r = x 1 ( w ) . If S indexes a subset of features in the transformed space indicating which part of the transformed input to interpret , we then define TRIM , f ( wS ) = attr ( f 0 ; wS ) , ( 7 ) where attr ( ; w ) is an attribution method that is evaluated at w and outputs an importance value , and where wS denotes the subvector of w indexed by S. The choice of the attribution method attr ( ) can be any local interpretation technique ( e.g . LIME [ 35 ] or Integrated Gradients ( IG ) [ 37 ] ) ; here we focus mainly on the saliency map [ 41 ] , which simply calculates the gradient of the model ’ s output with respect to its transformed input to 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 as capture information about a 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 reconstruction 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 and the corresponding scaling function span a 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 f for each coefficient in the wavelet representation , and try to find a wavelet function that makes these attributions sparse . Intuitively , this ensures that the learned wavelet should find a 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 a final optimization problem for AWD : minimize h , g L ( h , g ) = 1 m X i kxi bxik22 | { z } Reconstruction loss + 1 m X i W ( h , g , xi ; ) | { z } Wavelet loss + X i kTRIM , f ( xi ) k1 | { z } Interpretation loss , ( 8 ) where denotes a wavelet transform operator induced by , and bxi denotes the reconstruction of the data point xi . Here , > 0 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 , g ) via the following mapping [ 32 ] : b ( w ) = Q1 p=1 bh ( 2 pw ) p 2 and b ( w ) = 1p 2 bg ( w/2 ) b ( w/2 ) , where b and b represent the Fourier transforms of and 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 X n h [ n ] = p 2 and |bh ( w ) |2 + |bh ( w + ⇡ ) |2 = 2 for all w , ( 9 ) as well as some mild conditions , it can generate a scaling function such that the scaled and translated family of the scaling function forms an orthonormal basis of the space of multiresolution approximations of L2 ( 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. , X n h [ n ] h [ n 2k ] = ⇢ 1 if k = 0 0 otherwise , and as a result , khk2 = 1 . ( 10 ) 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 ) nh [ N 1 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.e. , P n g [ n ] = 0 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 ` 1 norm penalty on the wavelet coefficients . Combining all these constraints via regularization terms , we define the wavelet loss at the data point xi as W ( h , g , xi ; ) = k xik1 + ( X n h [ n ] p 2 ) 2 + ( X n g [ n ] ) 2 + ( khk22 1 ) 2 + X w ( |bh ( w ) |2 + |bh ( w + ⇡ ) |2 2 ) 2 + X k ( X n h [ n ] h [ n 2k ] 1k=0 ) 2 , where g is set as in ( 11 ) and > 0 controls strength of the sparsity of the wavelet representations . We enforce the penalty ( |bh ( w ) |2 + |bh ( w + ⇡ ) |2 2 ) 2 , only at the discrete values of w 2 { 2⇡kN , k = 1 , . . . , N } through the discrete Fourier transform . Notice that the wavelet loss does not introduce any additional hyperparameters besides . 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 . Interpretation loss The interpretation loss enables the distillation of knowledge from the pretrained model f into the wavelet model . It ensures that attributions in the space of wavelet coefficients xi are sparse , where the attributions of wavelet coefficients is 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 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 a 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 .

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.