Unveiling Transformer Perception by Exploring Input Manifolds

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 17:33:42

Abstract: This paper introduces a general method for the exploration of equivalence classes in the input space of
Transformer models. The proposed approach is based on sound mathematical theory which describes the internal layers of a Transformer architecture as sequential deformations of the input manifold. Using eigendecomposition of the pullback of the distance metric defined on the output space through the Jacobian of the model, we are able to reconstruct equivalence classes in the input space and navigate across them. We illustrate how this method can be used as a powerful tool for investigating how a Transformer sees the input space, facilitating local and task-agnostic explainability in Computer Vision and Natural Language Processing tasks.

Main paper: 
 introduction: 
 In this paper, we propose a method for exploring the input space of Transformer models by identifying equivalence classes with respect to their predictions. We define an equivalence class of a Transformer model as the set of vectors in the embedding space whose outcomes under the Transformer process are the same. The study of the input manifold on which the inverse image of models lies provides insights for both explainability and sensitivity analyses. Existing methods aiming at the exploration of the input space of Deep Neural Networks and Transformers either rely on perturbations of input data using heuristic or gradient-based criteria [16, 22, 17, 14] , or they analyze specific properties of the embedding space [5] .Our approach is based on sound mathematical theory which describes the internal layers of a Transformer architecture as sequential deformations of the input manifold. Using eigendecomposition of the pullback of the distance metric defined on the output space through the Jacobian of the model, we are able to reconstruct equivalence classes in the input space and navigate across them. In the XAI scenario, our framework can facilitate local and task-agnostic explainability methods applicable to Computer Vision (CV) and Natural Language Processing (NLP) tasks, among others.In Section 2, we summarise the preliminaries of the mathematical foundations of our approach.In Section 3, we present our method for the exploration of equivalence classes in the input of the Transformer models. In Section 4, we perform a preliminary investigation of some applicability options of our method on textual and visual data. In Section 5, we discuss the relevant literature about embedding space exploration and feature importance. Finally, in Section 6, we give our concluding remarks 1 .In this Section, we provide the theoretical foundation of the proposed approach, namely the Geometric Deep Learning framework based on Riemannian Geometry [2] .A neural network is considered as a sequence of maps, the layers of the network, between manifolds, and the latter are the spaces where the input and the outputs of the layers belong to.Definition 1 (Neural Network). A neural network is a sequence of C 1 maps Λ i between manifolds of the form:EQUATIONWe call M 0 the input manifold and M n the output manifold. All the other manifolds of the sequence are called representation manifolds. The maps Λ i are the layers of the neural network. We denote withN (i) = Λ n • • • • • Λ i : M i → M nthe mapping from the i-th representation layer to the output layer.As an example, consider a shallow network with just one layer, the composition of a linear operator A • +b with a sigmoid function σ, where A ∈ R m×n and b ∈ R m : then, the input manifold M 0 and the output manifold M 1 shall be R n and R m , respectively, and the map Λ 1 (•) = σ(A • +b). We generalize this observation into the following definition.Definition 2 (Smooth layer). A map Λ i : M i-1 → M i is called a smooth layer if it is the restriction to M i-1 of a function Λ (i) (x) : R di-1 → R di of the formEQUATIONfor i = 1, • • • , n, x ∈ R di , b (i) ∈ R di and A (i) ∈ R di×di-1 , with F (i) : R di → R di a diffeomor- phism.Remark 1. Transformers implicitly apply for this framework, since their modules are smooth functions, such as fully connected layers, GeLU and sigmoid activations.Our aim is to transport the geometric information on the data lying in the output manifold to the input manifold: this allows us to obtain insight on how the network "sees" the input space, how it manipulates it for reaching its final conclusion. For fulfilling this objective, we need several tools from differential geometry. The first key ingredient is the notion of singular Riemannian metric, which has the intuitive meaning of a degenerate scalar product which changes point to point. M i = d i . ii) The maps Λ i are C 1 submersions. iii) Λ i (M i-1 ) = M i for every i = 1, • • • , n. iv)The manifold M n is equipped with the structure of Riemannian manifold, with metric g (n) . Definition 3 naturally leads to the definition of the pseudolength and of energy of a curve. P l(γ) = b a ∥ γ(s)∥ γ(s) ds = b a g γ(s) ( γ(s), γ(s))ds, E(γ) = b a ∥ γ(s)∥ 2 γ(s) ds (3)The notion of pseudolength leads naturally to define the distance between two points.Definition 5 (Pseudodistance). Let x, y ∈ M = R n . The pseudodistance between x and y is thenEQUATIONOne can observe that endowing the space R n with a singular Riemannian metric leads to have non trivial curves whose length is zero. A straightforward consequence is that there are distinct points whose pseudodistance is therefore zero: a natural equivalence relation arises, i.e. x ∼ y ⇔ P d(x, y) = 0, obtaining thus a metric space (R n / ∼, P d).The second crucial tool is the notion of pullback of a function. Let f be a function from R p to R q , and fix the coordinate systems x = (x 1 , . . . , x p ) and y = (y 1 , . . . , y q ) on R p and on R q , respectively.Moreover, we endow R q with the standard Euclidean metric g, whose associated matrix is the identity.The space R p can be equipped with the pullback metric f * g whose representation matrix reads asEQUATIONThe sequence (1) shows that a neural network can be considered simply as a function, a composition of maps: hence, takingf = Λ n • Λ n-1 • • • • • Λ 1 and supposing that M 0 = R p , M n = R q , thegeneralization of (5) applied to (1) provides with the pullback of a generic neural network.Hereafter, we consider in (1) the case M n = R q , equipped with the trivial metric g (n) = I q , i.e., the identity. Each manifold M i of the sequence ( 1) is equipped with a Riemannian singular metric, denoted with g (i) , obtained via the pullback of N (i) . The pseudolength of a curve γ on the i-th manifold, namely P l i (γ), is computed via the relative metric g (i) via (3).Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: In the abstract and introduction we claim that we present a method for the exploration of equivalence classes in the input space of Transformer models, which is analyzed in depth in Section 3. The mathematical theory we refer to is deepened in Section 2. 
general results: 
 We depict hereafter the theoretical bases of our approach. We denote with N i the submapΛ i •• • ••Λ n :M i → M n , and with N ≡ N 0 the map describing the action of the complete network. The starting point is to consider the pair (M i , P d i ): this is a pseudometric space, which can be turned into a full-fledged metric space M i / ∼ i by the metric identification x ∼ i y ⇔ P d i (x, y) = 0. The first result states that the length of a curve on the i-th manifold is preserved among the mapping on the subsequent manifolds.Proposition 1. Let γ : [0, 1] → M i be a piecewise C 1 curve. Let k ∈ {i, i + 1, • • • , n} and consider the curve γ k = Λ k • • • • • Λ i • γ on M k . Then P l i (γ) = P l k (γ k ).In particular this is true when k = n, i.e., the length of a curve is preserved in the last manifold. This result leads naturally to claim that if two points are in the same class of equivalence, then they are mapped into the same point under the action of the neural network.Proposition 2. If two points p, q ∈ M i are in the same class of equivalence, thenN i (p) = N i (q).The next step is to prove that the sets M i / ∼ i are actually smooth manifolds: to this aim, we introduce another equivalence relation: x ∼ Ni y if and only if there exists a piecewise γ : [0, 1] → M i such that γ(0) = x, γ(1) = y and N i • γ(s) = N i (x) ∀s ∈ [0, 1].M i /∼ i = M i /∼ Ni+1 . Moreover, if two points p, q ∈ M i are connected by a C 1 curve γ : [0, 1] → M i satisfying N i (p) = N i • γ(s)for every s ∈ [0, 1], then they lie in the same class of equivalence.Making use of the Godement's criterion, we are now able to prove that the set M i / ∼ i is a smooth manifold, together with its dimension.Proposition 4. M i ∼ i is a smooth manifold of dimension dim(N (M 0 )).This last achievement provides practical insights about the projection π i on the quotient space, that consists the building block of the algorithms used for recovering and exploring the equivalence classes of a neural network.Proposition 5. π i : M i → M i / ∼ i is a smooth fiber bundle, with Ker(dπ i ) = VM i , which is therefore an integrable distribution. VM i is the vertical bundle of M i . Every class of equivalence[p] is a path-connected submanifold of M i and coincide with the fiber of the bundle over the point p ∈ M i .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss the limitations and tradeoff given by numeric integration in Subsection 3.1 and theoretical assumptions are enumerated in Section 2. Computational efficiency of our algorithms is discussed in Subsection 3.1. We conducted experiments on 3 datasets only, one of which of small dimensions since our main focus is on the mathematical theory grounding the application of the method to Transformers, as stated at the beginning of Section 4. Other limitations are mentioned throughout Section 4, including the fact that our investigations in the human-readable scenario are at a preliminary stage. 
methodology: 
 The results depicted in Section 2.1 provide powerful tools for investigating how a neural network sees the input space starting from a point x. In particular we point out the following remarks: i) If two points x, y belonging to the input manifold M 0 are are such that x ∼ 0 y, then N (x) = N (y); ii)given a point p ∈ M n , the counterimage N -1 (p) is a smooth manifold, whose connected components are classes of equivalences in M 0 with respect to ∼ 0 . A necessary condition for two points x, y ∈ M 0 to be in the same class of equivalence is that N (x) = N (y); iii) any class of equivalence [x], x ∈ M 0 , is a maximal integral submanifold of VM 0 . The above observations directly provide with a strategy to build up the equivalence class of an input point x ∈ M 0 . Proposition 5 tells us that VM 0 is an integrable distribution, with dimension equal to the dimension of the kernel of g (0) : we can hence find dim(Ker(g (0) )) vector fields which are a base for the tangent space of M 0 . This means that we can compute the eigenvalue decomposition of g (0)x and consider the L linearly independent eigenvectors, namely {v l } l=1,...,L , associated to the null eigenvalue: these eigenvectors depend smoothly on the point, a fact that is not trivial when the matrix associated to the metric depends on several parameters [15] . We can build then all the null curves by randomly selecting one eigenvector ṽ ∈ {v l } and then reconstruct the curve along the direction ṽ from the starting point x. From a practical point of view, one is led to solve the Cauchy problem, a first order differential equation, with γ = ṽ and initial condition γ(0) = x.This whole procedure is coded in the Singular Metric Equivalence Class (SiMEC) and the Singular Metric Exploration (SiMExp) algorithms, whose general schemes are depicted in Algorithms 1 and 2.SiMEC reconstructs the class of equivalence of the input via the exploration of the input space by randomly selecting one of the eigenvectors related to the zero eigenvalue. On the opposite, in SiMExp, in order to move from a class of equivalence to another we consider the eigenvectors relative to the nonzero eigenvalues. This requires the slight difference in lines 5 to 7 between Algorithm 1 and Algorithm 2.The Singular Metric Equivalence Class (SiMEC) algorithm. 1: Set the network N ; choose the maximum number of iterations. Choose the input p0.2: for k = 0, 1, . . . , K -1 do 3: Compute g n N (p k )Compute the pullback metric g 0 p k 5: Diagonalize g 0 p k and find the eigenvectors {v l } l associated to the zero eigenvalue 6:Randomly select ṽ ∈ {v l } l 7:δ = 1/ max(eigenvalues of g 0 p k ) 8:p k+1 ← p k + δṽ 9: end for 10: Optionally: store {p k } k=0,...,K for optimizing future computations 11: Project p k to the nearest feasible region Algorithm 2 The Singular Metric Exploration (SiM-Exp) algorithm.1: Set the network N ; choose the maximum number of iterations. Choose the input p0.2: for k = 0, 1, . . . , K -1 do 3: Compute g n N (p k )Compute the pullback metric g 0 p k 5: Diagonalize g 0 p k and find the eigenvectors {w l } l associated to the non-zero eigenvalue 6:Randomly select w ∈ {w l } l 7: δ = 2/ max(eigenvalues of g 0 p k ) 8:p k+1 ← p k + δ w 9: end for 10: Optionally: store {p k } k=0,...,K for optimizing future computations 11: Project p k to the nearest feasible region There are some remarks to point out. From a numerical point of view, the diagonalization of the pullback may lead to have even negative eigenvalues: hence one may use the notion of energy of a curve, related to the pseudolength. The update rule for the new point (line 8) amounts to solve the differential problem via the Euler method: for a reliable solution, we suggest to choose a small step-length δ. On the other hand, if the value of δ is too small more iterations are needed to move away from the starting point sensibly. Therefore there is a trade-off between the reliability of the solution and the exploration pace. The proof of the well-posedness theorem for Cauchy problems, cf. [18, Theorem 2.1], yields some insights, suggesting to set δ equal to the inverse of the Lipschitz constant of the map N -which in practice we can estimate with the inverse of the square root of the largest eigenvalue λ M of the pullback metric g 0 p k . This is our default choice for Algorithm 1. We also note that Algorithm 1 is more sensitive to the choice of the parameter δ compared to Algorithm 2.To build points in the same equivalence class Algorithm 1 needs to follow a null curve closely with as little approximations as possible, namely with a small δ. In contrast Algorithm 2, whose goal is to change the equivalence class from one iteration to the next, does not have the same problem and larger δ are allowed. Out default choice is therefore to set δ = 2λ-1/2 Mfor Algorithm 2. As for the computational complexity of the two algorithms, the most demanding step is the computation of the eigenvalues and eigenvectors, which is O(n 3 ), with n the dimension of the square matrix g 0 p k [20] .Since all the other operations are either O(n) or O(n 2 ), we conclude that the complexity of both Algorithms 1 and 2 is O(n 3 ).Algorithms 1 and 2 allow for the exploration of the equivalence classes in the input space of a Transformer model. However, the points explored by these algorithms may not be directly interpretable by a human perspective. For instance, an image or a piece of text may need to be decoded to be "readable" by a human observer. Furthermore, we present an interpretation of the eigenvalues of the pullback metric which allows us to define a feature importance metric. We present two interpretability methods for Transformers based on input space exploration. Both methods are then demonstrated on a Vision Transformer (ViT) trained for digit classification [8] , and two BERT models, one trained for hate speech classification and the other trained for MLM [7, 19] . Input data x (image or text) 4: Retrieve segments x t = tT (x). 5: Choose segments P = {p|p ∈ x t } for updates; keep others unchanged. 6: Compute embeddings x e = eT (x t ). 7: Apply SiMEC or SiMExp on x e , updating embeddings for segments in P . 8: Outputs: Modified input embedding, one for each SiMEC/SiMExp iteration.Feature importance. Consider a Transformer model T whose architecture includes a tokenizer t T (or patcher for images) that segments the input so that each segment can be converted into a continuous representation by an embedding layer e T . This results in a matrix of dimensions n s × h, where n s represents the number of segments, and h denotes the hidden size of the model's embeddings. The eigenvalues of the pullback metric can be used to deduce the importance of each embedding and, by extension, the significance of the segments they represent, with respect to the final prediction. The process for determining the importance of textual tokens or image patches is outlined in Algorithm 3.The appearance of the resulting heatmaps varies according to the type of input used. An example of experiments with ViT on the MNIST dataset [12] is shown in Figure 1 that depicts heatmaps for two MNIST instances. Figure 2 , on the left, illustrates two experiment using Algorithm 3 on both a BERT model for hate speech detection and a BERT model for MLM.Interpretation of input space exploration. Using SiMEC and SiMExp to explore the embedding space reveals how Transformer models perceive equivalence among different data points. Specifically, these methodologies facilitate the sequential acquisition of embedding matrices p 0 . . . p K at each iteration, as detailed in Algorithms 1 and 2. Algorithm 4 implements a practical application of the [MASK] is predicted as "ham", with the most influential tokens being pizza and cheese.SiMEC/SiMExp approach with Transformer models. A key feature of this method is its ability to selectively update specific tokens (for text inputs) or patches (for image inputs) during each iteration. This selective updating allows us to explore targeted modifications that prompt the model to either categorize different inputs as the same class or recognize them as distinct. Unlike traditional approaches where modifications are predetermined, this method lets the model itself guide us to understand which data points belong to specific equivalence classes. To interpret embeddings resulted from the exploration process, they must be mapped back into a human-understandable form, such as text or images. The interpretation of an embedding vector depends on the operations performed by the Transformer's embedding module e T . If e T consists only of invertible operations, it is feasible to construct a layer that performs the inverse operation relative to e T . The output can then be visualized and directly interpreted by humans, allowing for a comparison with the original input to discern how differences in embeddings reflect differences in their representations (e.g., text, images). If the operations in e T are non-invertible, a trained decoder is required to reconstruct an interpretable output from each embedding matrix p 0 . . . p K . When using a BERT model, it is feasible to utilize layers that are specialized for the masked language modeling (MLM) task to map input embeddings back to tokens. This approach is effective whether the BERT model in question is specifically designed for MLM or for sentence classification. In the case of sentence classification models, it is necessary to select a corresponding MLM BERT model that shares the same internal architecture, including the number of layers and embedding size.Algorithm 5 depicts the process of interpreting Algorithm 4 outputs for both ViT and BERT experiments. After initializing the decoder according to the model type, the embeddings p 0 . . . p K need to be constrained to a feasible region. This region is defined by the distribution of embeddings derived from the original input instances. Next, the embeddings are decoded, and the selected segments for exploration are extracted. These segments are then used to replace the corresponding parts of the original input instance. Figure 3 depicts an example outcome of Algorithm 5 applied on a ViT exploration experiment. Given that the interpretation process includes both a capping step and a decoding step (lines 10 and 11 of Algorithm 5), it's important to note that there isn't a direct 1:1 correspondence between each iteration's update and the interpretation outcomes. Our primary focus is on exploring the input embedding space, rather than the input image or input sentence spaces.For further investigation, we provide a detailed discussion on considering interpretation outputs as alternative prompts in Section 4. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The full proofs are part of two previously published papers which we cannot disclose for anonymity requirements. We replicate the relevant proofs in the supplementary material, part of which will be removed from the final version of the paper, referencing to the other papers.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments: 
 Experiments are conducted on textual and visual data. We aim to perform a preliminary investigation of 3 features of our approach: (i) how the class probability changes on the decoded output of SiMEC/SiMExp, (ii) what is the trade-off between the quantity and the quality of the output, and (iii) how our method can be used to extract feature importance-based explanations.In the textual case, we experiment with hate speech classification datasets: we use HateXplain 2 [13] , which provides a ground truth for feature importance, plus a sample of 100 hate speech sentences generated by prompting ChatGPT 3 , which serve purposes (i) and (ii). In the visual case, we perform experiments on MNIST [12] dataset. For BERT experiments we generate prompts to inspect the probability distribution over the vocabulary for tokens updated by Algorithms 1 and 2. We decode the updated p 0 . . . p K using Algorithm 5, focusing on tokens updated through the iterations. For each of these decoded tokens, we extract the top-5 scores to obtain 5 alternative tokens to replace the original ones, creating 5 alternate prompts. We then extract the prediction i * = arg max i y i for the original sentence, which represents the output whose equivalence class we aim to explore. Finally, we classify the new prompts, obtaining the corresponding predictions Y = y (0) . . . y (K) , where each y (k) ∈ R N , N being the number of prediction classes. We visualize the prediction trend for the i * th value in every y (0) . . . y (K) categorizing the images into two subsets: those that lead to a change in prediction Input space exploration We measure the time required to explore the input space of a ViT with the SiMEC algorithm and compare it with a perturbation-based method. The perturbation-based method mimics a trial-and-error approach as it takes an input image and, at each iteration, perturbs it by a semi-random vector v t+1 = a t v t + ηϵ, where a t = 1 if y t = y t-1 , a t = -1 otherwise, ϵ is an orthogonal random vector from a standard normal distribution and η is the step length. With the perturbation, we obtain a new image, then check whether the model yields the same label for the new image. The perturbation vector is re-initialized at random from a normal distribution 20% of the times to allow for exploration. We construct this method to have a direct comparison with ours in the absence of a consolidated literature about the task.Y c = {y (k) ∈ Y | arg max i y (k) i ̸ = i * } and those that don't Y s = {y i ∈ Y | arg max i y (k) i = i * }.We train a ViT model having 4 layers and 4 heads per layer on the MNIST dataset 5 . The SiMEC algorithm is run for 1000 iterations, so that it can generate 1000 examples starting from a single image. In a sample of 100 images, the average time is approximately 339 seconds. 6 In the same time, the perturbation-based algorithm can produce up to 36000 images. However, we notice that the perturbation-based algorithm ends up producing monochrome (pixel color has zero variance) or totally noisy images, which provide little information about the behavior of the model. Excluding only the images with low color variance (< 0.01), we are left, on average, with 19 images (standard deviation 13.9). SiMEC, in contrast, doesn't present this behavior, as all 1000 images have high enough intensity variance and are thus useful for explainability purposes.As BERT has many more parameters with respect to our ViT model, processing textual data takes longer. Specifically, in a sample of 16 sentences, the average time needed to run 1000 iterations on a sentence is 7089 seconds, taking into account both MLM and classification experiments.Feature importance-based explanations We compare our method against Attention Rollout (AR) [1] and the Relevancy method proposed by Chefer et al. [6] . In the textual case, we provide a quantitative evaluation using the HateXplain dataset, which contains 20147 sentences (of which 1924 in the test set) annotated with normal, offensive and hate speech labels as well as the positions of words that support the label decision. We then measure the cosine similarity between the importance assigned by each method to each word in a sentence and the ground truth. Notice that, since the dataset contains multiple annotations, the ground truth y for each word w is obtained as the average of the binary labels assigned by each annotator, and therefore y(w) ∈ [0; 1]. We also normalize all scores in [0; 1] so to have them on the same scale. The average similarity achieved by our method is 0.707 (standard deviation σ = 0.302), against 0.7 (σ = 0.315) for Relevancy and 0.583 (σ = 0.318) for AR. This proves our method to be more effective in finding the most sensitive tokens for classification.We provide an example on image classification in the Supplementary Materials.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: Pseudo-code of the proposed algorithms is reported in Subsections 3.1 and 3.2 so to make the algorithms reproducible, plus our implementation is made available in the supplementary material. Experiments, including the complete setting, and the respective baselines are described in Section 4. 
related work: 
 Our work relates to embedding space exploration literature, and has at least one collateral applications in the XAI domain, namely producing feature importance-based explanations.Embedding space exploration. Works dealing with embedding space exploration mostly focus on the study of specific properties of the embedding space of Transformers, especially in NLP. For instance, Cai et al. [5] challenge the idea that the embedding space is inherently anisotropic [10] discovering local isotropy, and find low-dimensional manifold structures in the embedding space of GPT and BERT. Biś et al. [3] argue that the anisotropy of the embedding space derives from embeddings shifting in common directions during training. In the field of CV, Vilas et al. [21] map internal representations of a ViT onto the output class manifold, enabling the early identification of class-related patches and the computation of saliency maps on the input image for each layer and head. Applying Singular Value Decomposition to the Jacobian matrix of a ViT, Salman et al. [17] treat the input space as the union of two subspaces: one in which image embedding doesn't change, and another one for which it changes. Except for the last one, all the aforementioned approaches rely on data samples. By studying the inverse image of the model, instead, we can do away with data samples.Feature importance-based explanations. Feature importance is a measure of the contribution of each data feature to a model prediction. In the context of Computer Vision and Natural Language Processing, it amounts to giving a weight to pixels (or patches of pixels) in an image and tokens in a piece of text, respectively. In recent years, much research has focused on Transformers in both CV and NLP. Most approaches are based on the attention mechanism of the Transformer architecture. Abnar and Zuidema [1] quantify the overall attention of the output on the input by computing a linear combination of layer attentions (Attention Rollout) or applying a maximum flow algorithm (Attention Flow). To overcome the limitations [4] of attention-based methods, Hao et al. [11] use the concept of attribution, which is obtained by multiplying attention matrices by the integrated gradient of the model with respect to them. Chefer et al. [6] propose the Relevancy metric to generalize attribution to bi-modal and encoder-decoder architectures. Other methods are perturbation-based, where perturbations of input data are used to record any change in the output and draw a saliency map on the input. In order to overcome the main issue with such methods, i.e. the generation of outlier inputs, Englebert et al. [9] apply perturbations after the position encoding of the patches. In contrast with these methods, ours does not need arbitrary perturbations of inputs, and considers all parameters of the model, not only the attention query and key matrices. 
conclusions: 
 Our exploration of the Transformer architecture through a theoretical framework grounded in Riemannian Geometry led to the application of our two algorithms, SiMEC and SiMExp, for examining equivalence classes in the Transformers' input space. We demonstrated how the results of these exploration methods can be interpreted in a human-readable form and conducted preliminary investigations into their potential applications. Notably, our methods show promise for ranking feature importance and generating alternative prompts within the same or different equivalence classes.Future research directions include expanding our experimental results and delving deeper into the potential of our framework for controlled input generation within an equivalence class. This application holds significant promise for enhancing the explainability of Transformer models' decisions and for addressing issues related to bias and hallucinations. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [No] Justification: Although the impacts of XAI on society is broad and deep, in this paper we focus only on the technical problem of exploring the equivalence classes in the input space of Transformers, which doesn't add any specific impact to the discussion about XAI in general.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks. 
 
 Appendix: • The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example Justification: The standard deviation is reported for the experiment that supports the claims of the paper, i.e. the one on feature importance-based explanations, in Section 4. Standard deviation is also reported for the number of uninformative images produced by the perturbation-based baseline method in the Input space exploration experiment.• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).Justification: We comply with the terms of use of the datasets employed in the experiments, and we deem our work has no potentially harmful effect on people safety, security, discrimination, surveillance, harassment, nor on human rights. Our proposal does not contribute to spread bias and unfairness towards certain groups of people nor to harm the environment.• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.The code to reproduce our experiments can be found in the Supplementary Materials.Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS2024). Do not distribute.Model used: huggingface.co/ctoraman/hate-speech-bertUsing Adam optimizer, the model achieved the highest validation accuracy (96.25%) in 20 epochs.All experiments are based on the current PyTorch implementation of the algorithms and run on a Ubuntu 20.04 machine endowed with one NVIDIA A100 GPU and CUDA 12.4. 


Decision: The paper presents a method for exploring the equivalence classes of inputs of transformer models. Reviewers praised the clarity of exposition and mathematical rigor of the paper, which presents a theoretical framework based on differential geometry that is nevertheless readable for non-experts. The most significant weaknesses mentioned by the reviewers are the lack of quantitative results (b3eP), lack of comparisons to existing feature importance methods (b3eP, KhQb) and the use of simple datasets like MNIST (TeNK).

In terms of ratings, we have a weak accept and accept (both with low confidence (with EZRw stating they are not expert on the topic), a confident reject and a low confidence borderline reject. 

In my view, the theoretical content of the paper is solid and interesting, but the paper is not yet ready for publication. I would recommend that the authors focus on experimenting with more challenging datasets, since results on MNIST often don't transfer. It would also be good to establish quantitative metrics to evaluate the success of the method, and make sure to do a thorough review of the literature and compare to all relevant prior works.

Reviews: review: paper_summary: paper_summary: This paper develops a novel theoretical framework grounded in Riemannian geometry for analyzing the input space of Transformer models, and introduce two algorithms, SiMEC and SiMExp, which facilitate the exploration and interpretation of equivalence classes within this input space. These methods offer new insights the internal mechanisms of Transformers, and provide new understanding of how these models perceive and process input data which can be very useful in the field of explainable AI.
strength_weakness: strength_weakness: In experiment, the MNIST dataset is a little bit trivial, as the pixels of the background is essentially zero. It is nice to see the application of the proposed algorithm on natural images like CIFAR.
questions: questions: See above
limitations: limitations: Yes
paper_summary: paper_summary: The authors present a method for exploring equivalence classes in the input space of Transformer models using a solid mathematical theory. By analyzing the Jacobian of the model, the method reconstructs and navigates these classes, offering a powerful tool for understanding Transformer interpretations and enhancing explainability. The proposed method is expected to solve problems in Computer Vision and Natural Language Processing tasks.
strength_weakness: strength_weakness: See above.
questions: questions: See above.
limitations: limitations: N/A.
paper_summary: paper_summary: This work attempts to find the set of inputs that generate the same neural network predictions. To this end, the authors interpret the layers of the network as transformations of the input manifold. This interpretation is used to defined equivalence classes over the inputs and to define feature importance. Finally, the tools are used to identify equivalence classes for MNIST digits and for hate speech detection with BERT.
strength_weakness: strength_weakness: The main contribution of this work is to introduce a tool for analyzing which set of inputs produce the same output. However, this is exactly the Fisher information matrix (with respect to the inputs) and has also been introduced in prior work (https://arxiv.org/abs/2104.13289). Could the authors clarify what the differences are and what the additional novelty. If the "local data matrix" introduced in https://arxiv.org/abs/2104.13289 is identical to the tools in this work, I think it severely diminishes the contributions of this work. Furthermore the experiments are extremely similar (such as Figure 1). 

The second weakness is the limited number of experiments. The work does not show any quantitative results: Figure 1, Figure 2 and Figure 3 is just 1 example and is not indicative of why the tools are useful. The experiments in section 4 primarily discuss wall-clock time. It would significantly help if claims such as "(Line 266) we notice that the perturbation-based algorithm ends up producing monochrome ..." are substantiated quantitatively. Overall, the work doesn't provide novel tools and the experiments lack a novel usage of these tools and do not reveal any new insights.
questions: questions: 1. How does this work differ from https://arxiv.org/abs/2104.13289?  
2. What is algorithm 4 (exploration) used for? 
3. Many experimental details are lacking. It would help to add appendices clarifying what networks were used and how they were trained. Overall, the work feels incomplete in many respects
4. For feature importance-based explanations, there are many other methods like LIME/SHAP (and many more follow up papers) which the authors do not compare to. Why were they omitted?
5. Does this method scale to larger datasets, considering we have to compute outer-product of the gradients and compute the Eigenvalues of this matrix?
limitations: limitations: The authors address limitations of their work but it can be expanded upon. For example, the authors can discuss the time required to compute Eigenvalues, and other limitations such as not having any Eigenvalues to be 0 in Algorithms 3 / 4. Furthermore, their algorithm should work (in theory) for infinitesimal steps in the input manifold.
paper_summary: paper_summary: The authors propose a mathematically rigorous methodology based on Riemannian geometry for attributing network importance of tokens in a transformer models input space (e.g. image patches, or ~words in the textual domain). The proposed methodology—whilst based on sound theory—translates into an intuitive algorithm involving what appears to be a relatively inexpensive eigendecomposition. Experiments on 3 datasets across both the image and NLP domains explore how the features correlate with ground-truth inputs in the text domain, in addition to first steps towards exploring how the features affect the networks’ output logits.
strength_weakness: strength_weakness: # [W1] Feature importance comparisons

Feature importance-based explanations are motivated on [L303] as quantifying the contribution of features `"to a model prediction"`. More concretely, around [L180], the authors motivate the eigenvalues of the pullback metric found using their method as ultimately deducing the importance of each segment (e.g. image patch)  `“with respect to the the final prediction”`. 

Consequently, a major weakness of the paper is how there is no comparison with related work for how well the proposed method’s identified important features alter the **output** logits (e.g. upon ablation).

I am slightly confused by why the authors did not adopt the established “perturbation test” experimental protocol in the baseline [1] against which they compared, to provide experimental evidence in favor of this. Currently, the only comparisons made around [L274] measure the features’ importance as they correlate to the *input’s* labels.

Concretely, the authors could, for example, ablate particular patches of MNIST and observe that the resulting performance drops correlate with the pullback metric’s eigenvalues. This would provide stronger evidence of the authors’ claims about the features affecting the networks’ output, and (crucially) ground the results in contrast to those achievable by existing methods.

# [W2] Limited experimental results & improvements

There is a lot of interesting theory here, but ultimately this is a paper with a concrete applied goal of feature attribution in transformer models. With such a new methodology with many technical details, I believe there is an extra burden of proof on the authors to demonstrate this somehow leads to additional insights / practical gains. As such, it is a relative weakness of the paper that so few experiments are performed to justify the methodology.

Beyond toy datasets, it would be interesting to see how the method performs on more complex ones (not necessarily larger ones), such as TinyImageNET. Here, we could visualize much more easily if the method helps identify salient features of animals’ body parts (for example) as being important features for classification. MNIST experiments alone in the image domain are hard to interpret given the similarity of all the input data. 

Furthermore, the method provides an almost insignificant increase of just `0.07` cosine similarity (over the baseline in [1]), on just a single dataset (and with just two baselines—for example, how does GradCAM perform here?). This is not sufficient evidence to convince me as a reader that the proposed methodology should be adopted. 

---

- [1]: Chefer, Hila et al. “Generic Attention-model Explainability for Interpreting Bi-Modal and Encoder-Decoder Transformers.” 2021 IEEE/CVF International Conference on Computer Vision (ICCV) (2021): 387-396.
questions: questions: No additional questions at this time (beyond those alluded to in the weaknesses section).
limitations: limitations: Some limitations are indeed addressed throughout. However, (unless I have missed something, in which case I apologise!) I can only find the limitations of the small number of datasets used stated in the NeurIPS checklist. This needs to be stated explicitly in the main paper.

Comments: title: Thank you for the response
Thanks to the authors
comment: We thank the reviewer for their effort in replying to our rebuttal response. However, we still do not agree with the reviewer’s point of view about the novelty of both our mathematical framework and algorithms. We are confident in the distinctiveness of our work compared to https://arxiv.org/pdf/2104.13289. In addressing the reviewers' concerns about its novelty, we thoroughly examined both the mathematical foundations and the application to Transformer architecture. Our responses highlighted six key differences and clarified our methodological contributions, emphasizing the relevance of all presented algorithms. While there may be surface-level similarities, our algorithms are underpinned by unique mathematical proofs that ensure broader generalizability (without restrictions on loss functions and covering a wider range of activation functions), applicability to trained networks (and not only on partially trained ones), and the flexibility to operate without assuming a constant rank for the pullback metric. Based on our comparison, we believe there are substantial differences between the two works, and the divergence points that we highlighted provide strong evidences against the reviewer’s argumentation.
Thank you for the taking the time to respond to all the questions. However many of my concerns remain unresolved. I agree with the authors that there are some differences with prior work that uses the local Fisher/data matrix. However, the differences about the activation function or loss function seem of little relevance in the context of deep networks. Unless I am mistaken, the algorithm used in (https://arxiv.org/pdf/2104.13289) is identical to algorithm 1. 

As mentioned in the weakness section, the experiments are limited to small datasets and 3 of the 4 results are qualitative. The purpose of algorithm 4/5 is still not clear to me; I understand the steps of the algorithm, but I don't understand what goal it achieves and there are no results that show it can be used to derive new insights. The authors have introduced new tools / algorithms but I believe that there needs to be more thorough experimental validation.
Thanks to the authors for their response.

I appreciate the authors’ efforts to clarify the goals and claims of the paper in their first response.

Unfortunately, I am still not convinced as it stands that the limited experimental results on MNIST sufficiently demonstrate the authors’ claim made in the rebuttal to `"demonstrate the potential empirical applications of our approach"` . I do not agree with the authors that adding additional results would have traded-off clarity (these can always be deferred to the supplementary material, for example).

Furthermore, two additional reviewers have the same flavor of concerns—it would have been easy for the authors to conduct small experiments on any of the additional suggested datasets (such as CIFAR by **Reviewer-TeNK** or TinyImageNET as suggested in the initial review). In the absence of additional experiments, my rating currently remains the same.