{'fraction_non_alphanumeric': 0.03208772182946184, 'fraction_numerical': 0.011196075602366181, 'mean_word_length': 4.709390444810544, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Locally Aggregated Feature Attribution on Natural Language Model Understanding

With the growing popularity of deep-learning models, model understanding becomes more important. Much effort has been devoted to demystify deep neural networks for better explainability. Some feature attribution methods have shown promising results in computer vision, especially the gradient-based methods where effectively smoothing the gradients with reference data is the key to a robust and faithful result. However, direct application of these gradient-based methods to NLP tasks is not trivial due to the fact that the input consists of discrete tokens and the “reference” tokens are not explicitly defined. In this work, we propose Locally Aggregated Feature Attribution (LAFA), a novel gradient-based feature attribution method for NLP models. Instead of relying on obscure reference tokens, it smooths gradients by aggregating similar reference texts derived from language model embeddings. For evaluation purpose, we also design experiments on different NLP tasks including Entity Recognition and Sentiment Analysis on public datasets and key words detection on constructed Amazon catalogue dataset. The superior performance of the proposed method is demonstrated through experiments.


Introduction
With the growing popularity of deep-learning models, model understanding becomes more and more critical in many folds. In one aspect, model understanding helps us understand what the model is doing by identifying crucial features among unstructured raw data. For example, Shrikumar et al., 2017 utilized the model explainability technique to discover motifs in regulatory DNA elements from distinct molecular signatures in the field of Genomics. In another aspect, model understanding helps people audit or debug the deep models. An interesting example is that Ribeiro et al. (Ribeiro et al., 2016) found that their image classification model sometimes misclassifies a husky as a wolf.
The model explainability tool reveals that their model relies on the snow in the background rather than the appearance when distinguishing the two animals. More importantly, model understanding helps gain trust when making important decisions based on the model. In the NLP domain, deep language models are quickly evolving and show superior performance in various benchmark tasks. However, even experts struggle to understand the mechanism of complex language models.
Much effort has been devoted to demystifying the "black box" of deep models. A natural idea is through feature attribution, explaining the model by attributing the prediction to each input feature according to how much it affects the model output, of which two main directions emerge. One is model agnostic approaches including Shapley regression values (Shapley, 1953) and LIME (Ribeiro et al., 2016). We can apply these methods regardless of the model structure, however, they could suffer from computational inefficiency in the scenario of high dimensional input space and complex deep models when making inferences across all possible permutations or with small perturbations in the local neighborhood.
Another direction is model-specific approaches which look into the internal model mechanism to understand specific models. Gradient-based feature attribution models are often adopted to explain neural networks since gradients can be easily accessed through back-propagation, which gives a great computational advantage over model-agnostic methods. Since the gradient map itself is often noisy and challenging to interpret, most gradient-based methods aim to stabilize the feature attribution score by smoothing the gradients or learning from the reference data (Sundararajan et al., 2017;Smilkov et al., 2017;Lundberg and Lee, 2017). However, direct application of these gradient-based methods to NLP problems is not trivial, due to the fact that the input consists of discrete tokens and the "reference" tokens are not explicitly defined.
In this paper, we propose Locally Aggregated Feature Attribution (LAFA), a novel gradientbased approach that leverages sentence-level embedding as a smoothing space for the gradients, motivated by the observation that the feature attribution is often shared by similar text inputs. For example, key features in product descriptions on an online marketplace are often shared by similar products. We implement a neighbor-searching method to ensure the quality of neighboring sentences.
Furthermore, to evaluate feature attribution methods in NLP, we consider two situations. For datasets with golden labels of feature score, we use the Area Under Curve (AUC) or Pearson correlation as the performance metric. As for datasets without golden labels, we conduct a similar evaluation task following prior works (Shrikumar et al., 2017;Lundberg and Lee, 2017) by masking tokens with high importance scores and find the change in the predicted log-odds.
In summary, our contributions are threefold: First, we build a novel context-level smooth gradient approach for feature attribution in NLP. The key ingredients of our method are constructing an appropriate aggregation function over the smoothing space. Second, to the best of our knowledge, this is the first proposal to conduct numerical studies on multiple NLP tasks, including Entity Recognition and Sentiment Analysis, for feature attribution. Third, our method achieves superior performance compared with the state-of-the-art feature attribution methods.
The paper is organized as follows. Section 2 elaborates the current challenges of feature attribution in NLP and recaps the preliminaries about gradient-based feature attribution approaches. The proposed feature attribution method is described in section 3, followed by a review of other existing approaches in Section 4. The evaluation tasks and the application results on NLP are presented in Section 5.

Feature Attribution in NLP
Challenge Direct application of gradient-based methods to NLP problems is not trivia. There are three main challenges. First, NLP models consist of non-differentiable discrete input tokens, hence the gradient hook can only reach out to the embedding space and gradient-based feature attribution methods are not directly applicable to word tokens.
Second, the reference data in NLP are difficult to define. It is studied by Sundararajan. et al (Sundararajan et al., 2017) that using the gradient as the feature attribution may suffer from the problems of model saturation or thresholding. Model saturation means the perturbation of some elements in the input cannot change the output, and the thresholding problem indicates discontinuous gradients can produce misleading importance scores. Such problems can be addressed by comparing the difference between the gradient of input and reference data. The guiding principle to select reference data is to ask ourselves that "what am I interested in measuring differences against?" For example, in the tasks of binary classification on DNA sequence inputs, the reference data are chosen as the expected frequencies of DNA sequence or randomly shuffling the original sequence. However, in NLP tasks, randomly shuffling texts as reference may not be grammatically sensible.
Lastly, we note that the evaluation of the language model is much more challenging than the explanations of the images. In the image application, the important features of an image obtained from feature attribution methods can be visually validated by checking the composition of objects. However, the detected important features in language may require more domain knowledge to validate.
Problem Definition Feature attribution task can be formally formulated as follows. A deep model F is provided to be explained, which is fine-tuned on dataset X . The input sentence is denoted as X 0 = (w 1 , w 2 , .., w T ) T where w i represents i-th word. The goal for feature attribution is to determine function M(·) by quantifying the importance score of each word M(x) = (m 1 , m 2 , .., m T ) T , where m i denotes the importance score for w i .
Simple Gradient as Feature Attribution As illustrated in the first challenge above, in NLP models, directly taking derivative on each word is infeasible due to the non-differentiable embedding layer. We can resolve the challenge as follows.
The fist layer of the NLP model usually maps input discrete tokens to embedding from a predefined dictionary.
where h 0,i ∈ R d represents the word embedding for w i . This step is non-derivative. But we can Input sentence of Interest: Dataset: Step 1 Localization … Similar texts: Step 2 Gradient … Step 3 Aggregation Output Step 1 Step 2 Deep Model: Step 3 Figure 1: Upper panel shows the overview of LAFA methods; Lower panel provides a motivating example of LAFA method. In this motivating example, the input text is a description of computer. The key features of the computer should include "Brand", "CPU type" and "RAM size". The simple gradient method may miss certain feature, such as "RAM size" in the example, while the gradients on similar texts can provide more contexts. The proposed method is constructed to aggregate the information from similar texts summarized in Algorithm 1.
obtain the derivative of output with respect to the word embedding: where H 0 = (h 0,1 , h 0,2 , .., h 0,T ) T ∈ R T ×d . Then, we consider the feature attribution score of a token M(X) ∈ R T as the sum of squares of the gradients with regard to each word embedding dimension: However, simply using the gradients of one token as feature attribution would lead to noisy results (Sundararajan et al., 2017). The next section describes a novel feature attribution approach that smoothes the gradients by leveraging similar input texts.

The Proposed Framework: LAFA
The proposed method contains three steps: (1) find the appropriate neighbors of the input text for gradient smoothing; (2) calculate the gradients of texts as well as neighbors; (3) aggregation of the gradients. The proposed framework is summarized in the upper panel of Figure 1. One motivating example is shown in the lower panel of Figure   1. In this motivating example, the input text is a description of computer. The key features of the computer should include "Brand", "CPU type" and "RAM size". The simple gradient method may miss certain feature, such as "RAM size" in the example, while the gradients on similar texts can provide more contexts. The proposed method is constructed to aggregate the information from similar texts.
Step I: Context-level Localization Given the input text X 0 ∈ X , where X denotes the input datasets, the goal is to find similar texts X sim = {X 1 , X 2 , .., X M } ⊂ X such that the feature attributions of X 0 and X j ∈ X sim are similar under a defined similarity metric.
To obtain similar texts X sim , we first define an encoder that maps the text with discrete word tokens to a continuous embedding vector; then, in the embedding space, similar texts are found in the neighbor of X 0 . To be specific, let H encoder denote the mapping from input to one of the hidden layers in deep model F. X sim can be obtained by choosing closest texts in the dataset as follows: where || · || 2 represents L 2 norm. ϵ is a threshold score to guarantee that founded neighbors are simi-lar to the center text X 0 to improve the faithfulness of aggregation. In our application, a fixed quantile served as the cut-off rate of L2 distance for all possible pairs is chosen as the threshold score to filter the nearest-neighbor result. During inference time, we apply the hidden layer encoder H encoder to all the input datasets and index, then using FAISS 1 (Johnson et al., 2017) offline. FAISS is an efficient, open-source library for similarity search and clustering on dense vectors, which can be applied to large-scale vectors. The output of this step, X sim can be viewed as the reference data to smooth the feature attribution of X 0 , which addresses the second challenge listed in Section 2.
Step II: Taking Gradients According to Equation (3), the gradient of X i can be denoted as To be noticed that our proposed method can be easily extended to variants of simple gradient including smooth gradient or integrated gradient methods (Smilkov et al., 2017;Sundararajan et al., 2017) in Step II.
Step III: Aggregation over Multiple Feature Attribution Our goal is to smooth the gradient M(X 0 ) by aggregating the gradients of similar text inputs: Since the lengths of X 0 , X 1 ,..,X M may vary, the lengths of gradients M(X 0 ), M(X 1 ),..,M(X M ) are different as well. Consequently, aggregation by simply taking the average is infeasible. Following the intuition that, the tokens with high gradients in X sim should be important in X 0 , we propose the following aggregation function: where λ is a hyper-parameter for leveraging the feature attribution from similar inputs. E(w; X sim ) is a scalar representing the importance of token w obtained from the neighbor inputs X sim . Formally, it can be defined as where h, h i,k are the word embedding of w and w i,k as in Equation (1) respectively, and k(·, ·) is a kernel function (Hofmann et al., 2008) (examples of kernel function are listed in the Appendix E. ). According to Equation (7), if word w and w i,k have a high similarity, then inner product between the embeddings h and h i,k in the kernel space would be large, which would assign a large weight to the corresponding importance score m i,k . On the contrary, dissimilar word w i,k in X sim has little effect to the word w in E(w; X sim ). The whole process is summarized in Algorithm 1.
Algorithm 1 Feature attribution method with smoothing over similar inputs. 1: Input: Text of interest X 0 , input datasets X , and fine-tuned deep model F. 2: Output: Feature attribution of X 0 3: Step I: Localization 4: Construct encoder H which maps from input space to the space of hidden layer in F. 5: Obtain the similar texts set X sim = {X 1 , X 2 , .., X M } of X 0 according to Equation (4). 6: Step II: Taking Gradient 7: Calculate the gradient of texts X 0 , X 1 , .., X M according to Equation (3). 8: Step III: Aggregation 9: Smooth the gradient of text of interest over the gradient of similar texts according to Equation (6). 10: Output the aggregated gradient M LAF A (X 0 ) as the feature attribution.
Discussion of Faithfulness One important criteria for model explainability method is "faithfulness", which refers to how accurately it reflects the true reasoning process of the model (Jacovi and Goldberg, 2020). In our proposed method, the original input X 0 is infused with similar texts in the input dataset X for better interpretation. Since the deep model F is also trained on X , using similar texts X sim ⊂ X to facilitate explanation will not violate the faithfulness.
In the localization step (Step I), out of the consideration about faithfulness, we do not use popular bi-encoder frameworks, such as S-BERT (Reimers and Gurevych, 2019) or DenseRetrival (Karpukhin et al., 2020), to obtain similar neighbors. Because it will involve an extra black box model when explaining deep model F.

Related Work
In NLP, transformer-based models yield great successfulness and some works focus on explaining the attention mechanism. For example, Serrano and Smith, 2019 and Jain and Wallace, 2019 inspected a single attention layer and found out that attention weights only weakly and inconsistently correspond to feature importance; Wiegreffe and Pinter, 2019 argued that we cannot separate the attention layer and should view the entire model as a whole. In this section, We mainly review the gradient-based methods for feature attribution.
Feature Attribution on Single Input Simonyan et al (Simonyan et al., 2013) computed the "saliency map" denoted as Simple Gradient from the derivative of the output with respect to the input in an image classification task. In the NLP application, "saliency map" is obtained as the derivative of the output with respect to the word embedding as in Equation (2). However, "saliency map" can be visually noisy. Several methods are proposed to improve the gradient method from different perspectives. Gradient*Input method (Shrikumar et al., 2017) improves the visual sharpness of the "saliency map" by multiplying gradient with the input itself. In NLP, we can write it as: Layerwise Relevance Propagation method (Bach et al., 2015) is shown to be equivalent to the Gradient*Input method up to a scaling factor. Smooth Gradient method (Smilkov et al., 2017) smoothes the feature attribution score by adding random noises to the input and taking average of the gradients from noisy inputs, formally: Guided Backpropagation method (Springenberg et al., 2014) modifies the back-propagation to preserve negative gradients in the ReLU activation layer which also sharpens the "saliency map" visually. Other methods, such as Grad-CAM or Guided-CAM (Selvaraju et al., 2017), are applicable to specific architecture of neural networks in the field of computer vision. Since language models like BERT do not contain specific architecture utilized in Guided Backpropagation or Grad-CAM method, we ignore the mathematical formulation here.
Feature Attribution on Input with Reference Data Integrated Gradient method computes the feature score by integrating the gradients from single pre-determined reference input to the target input (Sundararajan et al., 2017). In computer vision problems, black image is usually considered as the reference data, and integrating gradients from the black image to the input image represents the feature attribution of the input image. In NLP problems, we can define the i-th element of feature attribution as: where H ′ denotes the embedding of reference text. SHAP-Gradient method which combines ideas from Integrated Gradient and Smooth Gradient into a single expected value equation (Lundberg and Lee, 2017) . To be specific, the feature attribution is defined from: where α k ∼ U (0, 1) denotes uniform distribution from zero to one, H k ∈ H ref denotes the embedding of reference text.
DeepLIFT (Shrikumar et al., 2017) assigns the feature score by comparing the difference of contribution between input and some reference inputs via gradient. As discussed in (Lundberg and Lee, 2017), DeepLIFT can be considered as an approximation of Shapley Value estimation. Specifically, as in the application of SHAP 2 , the feature attribution of SHAP-Deep as a variant of DeepLIFT is defined as:

Experiments
In this section, we compare the proposed method to the state-of-the-art feature attribution methods under different use cases. Language Model We fine-tuned BERT-base models to classify the relations for NYT10 and Webnlg datasets, respectively. We use the plain text as input X, and relations as multi-class label Y in the model fine-tuning. Since multiple relations may exist in single text, we use the Sigmoid activation in the output layer. Mean Square Error (MSE) is used as loss objective and Adam (Kingma and Ba, 2014) is adopted as the optimizer. The micro Precision, Recall and F1 results are reported in Table 1 with 0.5 threshold of output score. From the result, the F1 scores are high for both NYT10 and Webnlg dataset, hence we can apply feature attribution methods to the fine-tuned models and identify the important features in the text which help to classify the relations.  Evaluation Metric In datasets, NYT10 and Webnlg, the positions of entities in triples are provided. Therefore, we can constructed the golden feature attribution label as follow. For text X = (w 1 , w 2 , .., w T ) T and triple (s, r, o), where subject s = (w i , .., w j ) and object o = (w k , .., w s ) are words shown in the text from positions i to j and k to s, respectively. The gold labels of feature attribution for relation r is constructed as  and M gold (X) for the test dataset. AUC ranges from 0 to 1, higher AUC represents the the feature attribution result is closer to the gold feature attribution.

Main Results
The results of AUC under different methods are summarized in Table 2. The popular feature attribution methods are listed and compared. More introduction about the competitors can be found in Section 4. "Rand", as a baseline method, denotes that the feature score is randomly assigned, therefore, the AUC score is about 0.5. InputGrad method performs better than the Simple-Grad, showing the effeteness of Taylor approximation of layer-wise relevance propagation. "SHAP + Zero" means zero references are used in SHAP and "SHAP + Ref." means X sim is used as references. SHAP-based methods show low AUC values, because such methods aggregate the gradients of input and reference by simply taking average aggregation (see details in Section 4), which is not meaningful in NLP tasks. From the result, our method LAFA achieves a superior performance in Webnlg dataset and comparable performance in NYT dataset, which indicates that our feature attribution method can identify entities well.

Case II: Feature Attribution on Sentiment Analysis
Motivation The goal of the sentiment classification task is to classify a text into a sentiment categories such as positive or negative sentiment (Aghajanyan et al., 2021;Raffel et al., 2019;Jiang et al., 2020). In this use case, we hope to explain the deep sentiment classification model and obtain sentiment factors that drive the model to identify the sentiment.
Data The Stanford Sentiment Treebank (SST) (Socher et al., 2013) is a sentiment analysis dataset collected from English movie reviews (Pang and Lee, 2005). For all 9, 645 sentences in SST, Amazon Mechanical Turk labeled the sentiment for words/phrases/sentences yielded from the Stanford Parser (Manning et al., 2014) on a scale between 1 and 25. SST-2 is first introduced by GLUE (Wang et al., 2018), a famous multi-task benchmark and analysis platform for natrual language understanding, which took a subset from the SST and applied a two-way split (positive or negative) on sentence-level labels. Owing to the fact that the train/validation/test split are aligned between SST and SST-2, we can run gradient-based methods on the either one of them. Note that we are only working with the test split for both data sets, which contains 2210 and 1821 sentences respectively.

Language Model
We use a popular and publicly available Distill-BERT (Sanh et al., 2019) model which is fine-tuned on SST-2 3 . The accuracy of the Distill-BERT model on SST and SST-2 is 86.6% and 92.4% respectively.

Evaluation Metric
We extract word-level sentiments from the phrase structure tree (PTB) in SST dataset. We take an absolute value after centralization to yield the golden label M gold (X). Pearson correlation coefficient, ρ , is the evaluation metric for feature attribution M gold (X) and M · (X). The correlation ρ takes value from the range from −1 to 1, and a higher ρ means better feature attribution result.

Main Results
The main results of the correlation are summarized in Table 3. Popular feature attribution methods are listed and compared. To leverage the problem that some words can have opposite meaning when their sentiment are different, we only limited the sentences neighbor for same category. Based on the preliminary experiment, we choose the second layer with 10 neighbors and 0.39 cut-off rate, more details about preliminary experiment can be found in Appendix D that using all layers of DistilBERT as the encoder will improve the performance. From the result Table 3, it is interesting to point out that the DistilBERT model fine-tuned on SST-2 does not perform equally well on the remaining sentences in SST, so the explanation we yield also has lower correlation for all methods compared with the SST-2. Some example and analysis when LAFA works and fails in this dataset by showing neighbor sentences can be found in Appendix B.
We can find out that the InputGrad method outperform SimpleGrad on SST/SST2 as well.
SmoothGrad method achieves a good result by introducing random noise. From our observation, Sharpley-Value based methods, "SHAP + Zero" and "SHAP + Ref." can identify important features with a good chance but may include several irrelevant tokens leading to higher variance. Our method LAFA achieves a superior performance in larger average correlation and smaller variance on both SST-2 and SST data sets.

Case III: Feature Attribution on Regression Model
Motivation Amazon's online stores contain rich information about millions of products in product title, brand and description. We hope to better understand the trendy features that affecting price directly from such unstructured raw data, without the need for human labelers / data cleaning. In this application, we fine-tuned a deep language model with price as labels and aim to understand important factors from product descriptions with the given language model.

Data
We collected the product catalog data of about one million products in personal computer category on Amazon's online store. We concatenate product's title, brand, bullet points and description as the input X, and use product price as the label Y .
Language Model We use BERT-base model and fine-tuned on collected catalog data for price regression.
Evaluation Metric To evaluate the performance of feature attribution methods without golden labels, we follow a similar idea as in work (Shrikumar et al., 2017;Lundberg and Lee, 2017) where the difference of prediction log-odds are measured by deleting pixels with highest importance scores. In our application, we first randomly select 200 input texts within a threshold of 1% prediction error as evaluation set. For each input text, we then mask p% of the tokens with highest feature attribution scores according to different feature attribution methods. Then we obtain new prediction result from the masked text denoted asŷ masked and calculate the new mean absolute percentage error (MAPE). Higher value of MAPE means that the corresponding method excels in picking important features.

Main Results
The results are shown in Figure  2 where x-axis is the mask proportion p, and yaxis is MAPE. We observe that the random method has very low MAPE, because randomly masking the input texts will not affect the predicted result as much as the other feature attribution methods. ShapDeep and ShapGrad also have low MAPE values since simply taking average as aggregation is meaningless in NLP tasks. Other competing methods have similar performances on this case study and non of these performs better than others in a wide range of mask ratio. The proposed LAFA method outperforms other methods by significant margin with masking proportion from 5% to 50%, which demonstrates that smoothing over contextlevel neighbors helps to highlight the important features in similar type of products.

Conclusion
This paper presents a novel locally aggregated feature attribution method in NLP, which efficiently captures the important features by leveraging similar input texts in the embedding space. We focused on feature attribution of single input based on a fine-tuned model instead of training a language model, henceforth the computation time is of less concern.
One limitation of the LAFA model is that it requires informative neighbor sentences that carry similar information. Otherwise, aggregating information from other sentences could be misleading. Experiments in our datasets show that our method is effective, but the improvements gained from the LAFA varies among different datasets based on the information that neighbor carries.
There are several future directions worthy of study. Firstly, labeling feature attribution result in the NLP requires massive human labor, and few datasets are available with golden feature attribution label. Developing new evaluation techniques to further measure model performance is interesting to investigate. Also, readable feature attribution results could help human beings to develop more business applications. For example, developing a key-value pair like processor-i5 as important feature can provide a more plausible feature attribution result to customers.

Appendix: A Model Implementation Detail
All experiments are conducted with eight NVIDIA Tesla V100 GPUs with 2.5 GHz (base) and 3.1 GHz (sustained all-core turbo) Intel Xeon 8175M processors. Case I For LAFA, we adopt the cosine function as the kernel function and hyper-parameter with λ = 1 and SimpleGrad is implemented and aggregated by M(x) in Equation (5). Case II For LAFA, we adopt the Polynomial function as the kernel function k(·, ·) = I(·, ·) and hyper-parameter with λ = 0.44 and SmoothGrad is chosen and aggregated by M(x) in Equation (5).
For gradient-based model with hyperparameters, we tuned them on the first 100 sentences in the test set. From the grid [10, 25, 50], we choose 25 as the integral iteration and the smooth candidates. Case III The indicator function as the kernel function k(·, ·) = I(·, ·) with λ = 1 as hyperparameter is adopted for LAFA. Neighbor information is aggregated by M(x) in Equation (5)  In Figure 3, we show two examples in NYT and Webnlg with their neighbors. We can observe that detected neighbor sentences have a similar meaning, which can be utilized as a reference to help extract the key features from the original sentence.

Example of Neighbors for NYT Dataset :
Example of Neighbors in Webnlg Dataset : Figure 3: Example of neighbors for NYT and Webnlg. The head and tail entities are highlighted with red color.

B.2 Sentiment Analysis
In SST-2, finding informative neighbors for every sentence is difficult because top sentences may not contain similar tokens, thus does not help. For this reason, we used a cut-off value for this data set to filter out non-informative sentences. In figure  4 we can find two examples from SST, one with "informative" good neighbors but another without them. Here for the word "informative" we use a quote because we are judging them based on our human understanding.

C Examples of Different Feature Attribution Methods under Multiple Cases
Here we provide an example in Cases I and II. In Figure 5, LAFA identified locations and "lived" as the important factors for relation extraction, and the importance of the "Atlantic City" and "Bader Field" is stronger than the backbone SimGrad because of aggregation.

D Experiment on Different Layer as Neighbor Encoder
Denote the size of X sim as M , the choice of which can be a critical and challenging task. Intuitively, an overly small M would lead to under-smoothing because the target text cannot incorporate enough information from the neighbors. On the contrary, an overly large M would cause over-smoothing by introducing too much noise. To clarify the neighbor searching process and the difference in the result using different layers, we show some experiments below.
Admittedly, we can directly use the WordPiece embedding as the encoder, which is the input of BERT-based models and enable us to find neighbors in the sense of "Word Similarity". However, since the same word can have different meanings in different sentences, and thus different importance in yielded gradients, we might need to use another layer in the BERT model as the encoder to incorporate contextual information.
We separate the layer search process into two cases depending on the availability of a set of labels that categorizes similar contents into the same group. Generally speaking, both cases recommend the middle layer as the encoder based on our experience.

D.1 When extra labels are not available
In the case of SST data, we do not have anything to group similar sentences, so we need to try for differ- Here we fixed the max number of neighbors as 10 and uses 0.05 quantile of sampled similarities as the cut-off rate to filter those neighbors that are not "actually close". We use the SimpleGrad and the SmoothGrad as the baseline for comparison on the first 100 sentences in the test set.
From table A1 we can find out that LAFA is a generally good method that always beats the baseline when we use the smooth gradient as the basement method. Layer 2 performs the best among candidates. The combination of SmoothGrad and Layer w is the final choice and we showed the results on entire SST in the main result part. From here we can find out that for all seven layers, information from faithful neighbors can bring some useful information to an existing sentence.

D.2 When we have extra label
The performance of encoders can be evaluated by the similarity between text X 0 and similar texts X sim obtained from Equation (4). In this application, we use the product category or subcategory which is an additional source of labels produced by Amazon to construct a proxy metric to evaluate the I(c(X j ) = c(X 0 )), (8) where c(·) denotes the category or subcategory of the corresponding product, I(·) is the indicator function. A high precision represents that the text found X sim are similar to the text of interest X 0 .
In the numerical study, we randomly sample 10, 000 inputs texts and obtain their corresponding neighbor texts from Equation (4) with M = 10 using each of the 12 hidden layers in BERT as the encoder H encoder under L 2 norm. Figure 7 shows the precision result from different encoders, where we observe that the fifth hidden layer has the highest precision in terms of both category and subcategory, which is consistent with the intuition that the middle layer is a trade-off of token-alike and output-alike inputs. In the following experiment, we adopt the fifth layer as the encoder. In general, when no external labels are provided, we may choose a different encoder depending on the use case.

E Ablation Study on Kernel Function
In Case III, we conduct an ablation study with different choices of kernel functions using different mask ratio to find out if different kernel yields different learning speed: 1. Radial basis function kernel (RBF) : k RBF (a, b) = exp(−||a − b|| 2 /l 2 ), where larger hyper-parameter l indicates lower impact from neighbors and vice versa. In the numerical study, we choose l = 2 based on the range of embedding a and b.
The results are shown in Figure 8. We observe that no single kernel function outperforms all other kernel functions under all mask ratios in this study. Indicator function shows a good performance when the masked ratio is greater than 10%, while RBF kernel shows a good performance when the masked ratio is smaller than 5%. This can due to the reason that the indicator function only aggregates identical words and this conservative manner helps when we lost most important words.