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 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. The 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.  The 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.