Original: Jump Self-attention: Capturing High-order Statistics in Transformers
Revised: Jump Self-attention: Capturing High-order Statistics in Transformers

Original: Abstract
Revised: Abstract

Original: The recent success of Transformer has benefited many real-world applications, with its capability of building long dependency through pairwise dot-products.
Revised: The recent success of Transformer has benefited many real-world applications, with its capability of building long dependency through pairwise dot-products.

Original: However, the strong assumption that elements are directly attentive to each other limits the performance of tasks with high-order dependencies such as natural language understanding and Image captioning.
Revised: However, the strong assumption that elements are directly attentive to each other limits the performance of tasks with high-order dependencies such as natural language understanding and Image captioning.

Original: To solve such problems, we are the first to define the Jump Self-attention (JAT) to build Transformers.
Revised: To solve such problems, we are the first to define the Jump Self-attention (JAT) to build Transformers.

Original: Inspired by the pieces moving of English Draughts, we introduce the spectral convolutional technique to calculate
Revised: Inspired by the pieces moving of English Draughts, we introduce the spectral convolutional technique to calculate

Original: JAT on the dot-product feature map.
Revised: JAT on the dot-product feature map.

Original: This technique allows JAT’s propagation in each self-attention head and is interchangeable with the canonical self-attention.
Revised: This technique allows JAT’s propagation in each self-attention head and is interchangeable with the canonical self-attention.

Original: We further develop the higher-order variants under the multi-hop assumption to increase the generality.
Revised: We further develop the higher-order variants under the multi-hop assumption to increase the generality.

Original: Moreover, the proposed architecture is compatible with the pre-trained models.
Revised: Moreover, the proposed architecture is compatible with the pre-trained models.

Original: With extensive experiments, we empirically show that our methods significantly increase the performance on ten different tasks.
Revised: With extensive experiments, we empirically show that our methods significantly increase the performance on ten different tasks.

Original:  Introduction
Revised:  Introduction

Original: Transformers have made great success in many domains, such as NLP, Vision, Time-series, and
Revised: Transformers have made great success in many domains, such as NLP, Vision, Time-series, and

Original: Speech.
Revised: Speech.

Original: However, the giant model GPT-3 [2] and switch Transformer [7], with their massive layers and over-parameterized architectures, have not further achieved significant performance improvements since the BERT model [5] beat the Human Performance in GLUE [21] and SQuAD [12].
Revised: However, the giant model GPT-3 [2] and switch Transformer [7], with their massive layers and over-parameterized architectures, have not further achieved significant performance improvements since the BERT model [5] beat the Human Performance in GLUE [21] and SQuAD [12].

Original: Recently two visualization works [19, 3] reveal that the self-attention mechanism mainly learns superficial connections among inputs.
Revised: Recently two visualization works [19, 3] reveal that the self-attention mechanism mainly learns superficial connections among inputs.

Original: The principle of direct accessing among inputs in the dotproduct self-attention lies as the core of above models, which solely relies on their mutual similarity in the vector space.
Revised: The principle of direct accessing among inputs in the dotproduct self-attention lies as the core of above models, which solely relies on their mutual similarity in the vector space.

Original: However, as “every coin has two sides”, while the dot-product is simple and effective in capturing long-range dependencies, it has limited capability in learning high-order dependencies.
Revised: However, as “every coin has two sides”, while the dot-product is simple and effective in capturing long-range dependencies, it has limited capability in learning high-order dependencies.

Original: To solve this problem, we are inspired by the pieces moving of English Draughts, a checker game spread to Great Britain in about 1100.
Revised: To solve this problem, we are inspired by the pieces moving of English Draughts, a checker game spread to Great Britain in about 1100.

Original: It allows two moving rules: one is the simple move from the
Revised: It allows two moving rules: one is the simple move from the

Original: [Figure]
Revised: [Figure]

Original: current position to an adjacent position; and the other is the capture move jumping over the “stepping stones” (opponent’s pieces) from front to back.
Revised: current position to an adjacent position; and the other is the capture move jumping over the “stepping stones” (opponent’s pieces) from front to back.

Original: Multiple jumps are possible and forced in this game.
Revised: Multiple jumps are possible and forced in this game.

Original: As shown in Fig.(1), we can employ the high-order self-attention connections, like the capture moves in English Draughts, to enhance the feature map of dot-product self-attention.
Revised: As shown in Fig.(1), we can employ the high-order self-attention connections, like the capture moves in English Draughts, to enhance the feature map of dot-product self-attention.

Original: It helps to align high-order global dependencies in Transformer models.
Revised: It helps to align high-order global dependencies in Transformer models.

Original: As a more concrete example in natural language understanding, we are given two sentences “Book costs more than pencil.”
Revised: As a more concrete example in natural language understanding, we are given two sentences “Book costs more than pencil.”

Original: and “Book costs less than computer”.
Revised: and “Book costs less than computer”.

Original: “Computer” never attends to “pencil” until they were both compared with “Book”, thus “Book” is the “stepping stones” in creating high-order attention between “computer” and “pencil” in Fig(1b).
Revised: “Computer” never attends to “pencil” until they were both compared with “Book”, thus “Book” is the “stepping stones” in creating high-order attention between “computer” and “pencil” in Fig(1b).

Original: More specifically, they forms a soft triadic closure.
Revised: More specifically, they forms a soft triadic closure.

Original: We will discuss this example in detail in Section 3.1 and clarify the definition of Jump Self-attention.
Revised: We will discuss this example in detail in Section 3.1 and clarify the definition of Jump Self-attention.

Original: In this work, we find that the high-order self-attention exists in the dot-product computation but it is too weak to contribute much to the attention layer’s output.
Revised: In this work, we find that the high-order self-attention exists in the dot-product computation but it is too weak to contribute much to the attention layer’s output.

Original: Only a few heads in Transformers can capture such high-order dependencies, which is shown in Fig.(2) with further discussions.
Revised: Only a few heads in Transformers can capture such high-order dependencies, which is shown in Fig.(2) with further discussions.

Original: Instead of finding such “miracle” heads, we reformulate them to enhance the high-order dependencies with the feature propagation technique in graph learning.
Revised: Instead of finding such “miracle” heads, we reformulate them to enhance the high-order dependencies with the feature propagation technique in graph learning.

Original: Meanwhile, we build an efficient variant in computing-sensitive situations based on the sparsity measurement [24], and build higher-order variants in comprehensive situations using MixHop [1].
Revised: Meanwhile, we build an efficient variant in computing-sensitive situations based on the sparsity measurement [25], and build higher-order variants in comprehensive situations using MixHop [1].

Original: To the best of our knowledge, we are the first to define the Jump Self-attention (JAT) in Transformer models, where we formulate the high-order statistics calculation into a graph propagation problem.
Revised: To the best of our knowledge, we are the first to define the Jump Self-attention (JAT) in Transformer models, where we formulate the high-order statistics calculation into a graph propagation problem.

Original: It can be used interchangeably with the canonical self-attention [18].
Revised: It can be used interchangeably with the canonical self-attention [18].

Original: The JAT-enhanced Transformer allows for fine-tuning on the popular pre-trained models and achieves superior performance on various tasks, such as Natural Language Understanding and Question Answering.
Revised: The JAT-enhanced Transformer allows for fine-tuning on the popular pre-trained models and achieves superior performance on various tasks, such as Natural Language Understanding and Question Answering.

Original: Our main contributions can be summarized as:
Revised: Our main contributions can be summarized as:

Original: • We discover the existence of JAT in the Transformers’ dot-product computation and derive its formal definition to facilitate further analysis.
Revised: • We discover the existence of JAT in the Transformers’ dot-product computation and derive its formal definition to facilitate further analysis.

Original: • We propose JAT through the jump aggregation and high-order dependencies’ propagation without additional network parameters, and the efficient variant alleviates the extra computation cost.
Revised: • We propose JAT through the jump aggregation and high-order dependencies’ propagation without additional network parameters, and the efficient variant alleviates the extra computation cost.

Original: • We propose the higher-order variants of JAT and masking strategy to allow its interchangeable usage with the canonical self-attention on the pre-trained models.
Revised: • We propose the higher-order variants of JAT and masking strategy to allow its interchangeable usage with the canonical self-attention on the pre-trained models.

Original: • We conduct experiments on fifteen different tasks to examine the effectiveness of the proposed methods.
Revised: • We conduct experiments on fifteen different tasks to examine the effectiveness of the proposed methods.

Original: Results demonstrate the success of JAT in achieving state-of-the-art performance.
Revised: Results demonstrate the success of JAT in achieving state-of-the-art performance.

Original:  Preliminary
Revised:  Preliminary

Original: The self-attention was introduced by Transformer [18] model for the encoder-decoder architecture in NLP to build long-range dependencies from a variable-length source sentence.
Revised: The self-attention was introduced by Transformer [18] model for the encoder-decoder architecture in NLP to build long-range dependencies from a variable-length source sentence.

Original: The alignment ability of self-attention to learn to acquire direct access between important items has made it an efficient component in many applications [5, 11, 24].
Revised: The alignment ability of self-attention to learn to acquire direct access between important items has made it an efficient component in many applications [5, 11, 25, 23].

Original: For an input X ∈ R L × d , we define the single-headed self-attention in the following form:
Revised: For an input X ∈ R L × d , we define the single-headed self-attention in the following form:

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: first they where the layer output is the linear combination of values V weighted by the attention scores between queries Q and keys K .
Revised: first they where the layer output is the linear combination of values V weighted by the attention scores between queries Q and keys K .

Original: Note that the queries, keys and values are projected by different linear transformations { W Q , W K , W V } respectively.
Revised: Note that the queries, keys and values are projected by different linear transformations { W Q , W K , W V } respectively.

Original: In practice, multiple attention heads are used to learn robust patterns [3, 19] among different transformations.
Revised: In practice, multiple attention heads are used to learn robust patterns [3, 19] among different transformations.

Original: It divides the inputs into H groups and computes the single-head attention on each group separately.
Revised: It divides the inputs into H groups and computes the single-head attention on each group separately.

Original: We concatenate the feature map of all groups as the attention layer output.
Revised: We concatenate the feature map of all groups as the attention layer output.

Original:  Methods
Revised:  Methods

Original: In this section, we first introduce the necessity of acquiring the high-order connections in the
Revised: In this section, we first introduce the necessity of acquiring the high-order connections in the

Original: Transformer model.
Revised: Transformer model.

Original: Then, we describe the proposed Jump Self-attention.
Revised: Then, we describe the proposed Jump Self-attention.

Original: 3.1 What is the High-order Statistics ?
Revised: 3.1 What is the High-order Statistics ?

Original: The BERT [5] model and its pre-trained successors utilize the self-attention, like GPT-3 [2] and switch
Revised: The BERT [5] model and its pre-trained successors utilize the self-attention, like GPT-3 [2] and switch

Original: Transformer [7], and they achieve distinct performance improvements on various NLP tasks.
Revised: Transformer [7], and they achieve distinct performance improvements on various NLP tasks.

Original: However, what the self-attention mechanism has learned remains to be a challenging problem for architecture optimization in Transformer.
Revised: However, what the self-attention mechanism has learned remains to be a challenging problem for architecture optimization in Transformer.

Original: Recent visualization techniques on self-attention mechanism have shed light on the further explanation for this problem.
Revised: Recent visualization techniques on self-attention mechanism have shed light on the further explanation for this problem.

Original: In the BERT model, Vig [19] pointed out that different heads of the same layer highlight coarse positional patterns, such as ‘CLS’ or ‘SEP’ tokens, more frequently than lexical patterns, even though they are trained with separate random projections.
Revised: In the BERT model, Vig [19] pointed out that different heads of the same layer highlight coarse positional patterns, such as ‘CLS’ or ‘SEP’ tokens, more frequently than lexical patterns, even though they are trained with separate random projections.

Original: Clark et al.
Revised: Clark et al.

Original: [3] claimed that BERT’s attention heads often exhibit similar behaviors within the same layer, such as attending to delimiter and positional tokens and broadly attending the entire sentence, whereas a few deviating heads hold an eye on linguistic pairs.
Revised: [3] claimed that BERT’s attention heads often exhibit similar behaviors within the same layer, such as attending to delimiter and positional tokens and broadly attending the entire sentence, whereas a few deviating heads hold an eye on linguistic pairs.

Original: Due to the space limit, a more comprehensive related work will be provided in the Appendix A.
Revised: Due to the space limit, a more comprehensive related work will be provided in the Appendix A.

Original: This paper investigates whether it is worth enhancing lexical patterns in self-attention.
Revised: This paper investigates whether it is worth enhancing lexical patterns in self-attention.

Original: To motivate our approach, we perform an individual case study on investigating the lexical patterns of self-attention heads.
Revised: To motivate our approach, we perform an individual case study on investigating the lexical patterns of self-attention heads.

Original: Based on the pre-trained model on BERT [5], we consider logical reasoning questions on two continuous sentences “Book cost more than pencil” and “Book cost less than computer”.
Revised: Based on the pre-trained model on BERT [5], we consider logical reasoning questions on two continuous sentences “Book cost more than pencil” and “Book cost less than computer”.

Original: It seeks to find the price between entities that have not been compared, like “pencil” and “computer”.
Revised: It seeks to find the price between entities that have not been compared, like “pencil” and “computer”.

Original: Note that the articles are omitted to avoid interference.
Revised: Note that the articles are omitted to avoid interference.

Original: In Fig.(2a), we present the visualization of three words’ self-attention scores at the third layer and fifth head.
Revised: In Fig.(2a), we present the visualization of three words’ self-attention scores at the third layer and fifth head.

Original: The darker line color denotes a higher self-attention score, namely stronger connections.
Revised: The darker line color denotes a higher self-attention score, namely stronger connections.

Original: The word “book” shows strong attention connections (deeper blue lines) to the positional token ‘SEP’ in the first and second sentences, meanwhile, they also attend to the entity words “pencil” and “computer” (lighter blue lines).
Revised: The word “book” shows strong attention connections (deeper blue lines) to the positional token ‘SEP’ in the first and second sentences, meanwhile, they also attend to the entity words “pencil” and “computer” (lighter blue lines).

Original: While the attention scores between the word “computer” and token ‘SEP’ are striking, we note that the word “computer” builds a mild attention connection to the word “pencil” crossing two sentences.
Revised: While the attention scores between the word “computer” and token ‘SEP’ are striking, we note that the word “computer” builds a mild attention connection to the word “pencil” crossing two sentences.

Original: From the entity extraction perspective, the mutual attention connections among three words “book”, “pencil” and “computer” seem trivial.
Revised: From the entity extraction perspective, the mutual attention connections among three words “book”, “pencil” and “computer” seem trivial.

Original: However, the concrete finding in Fig.(2b) undermines this emotional thinking and it reveals the potential value of attention connections between “computer” and “pencil”.
Revised: However, the concrete finding in Fig.(2b) undermines this emotional thinking and it reveals the potential value of attention connections between “computer” and “pencil”.

Original: We first calculate the mean value and standard deviation over all the pairwise self-attention scores.
Revised: We first calculate the mean value and standard deviation over all the pairwise self-attention scores.

Original: And we select the attention connections that their scores are greater than the threshold (mean + 0.5 std).
Revised: And we select the attention connections that their scores are greater than the threshold (mean + 0.5 std).

Original: We only draw the local connections among the three words and it covers 15.2% of the total as the significant local connections.
Revised: We only draw the local connections among the three words and it covers 15.2% of the total as the significant local connections.

Original: The stronger “Computer-Pencil” attention connection, as it in Fig(2b), is the rare self-attention computation that happens at the third layer, which distinguishes it from the trivial “Book-pencil” and “Book-Computer” attention connections.
Revised: The stronger “Computer-Pencil” attention connection, as it in Fig(2b), is the rare self-attention computation that happens at the third layer, which distinguishes it from the trivial “Book-pencil” and “Book-Computer” attention connections.

Original: As we discussed above, there are many connections from three words to the delimiter ‘CLS’ and positional tokens ‘SEP’ but they did not form the “Computer-Pencil” connection.
Revised: As we discussed above, there are many connections from three words to the delimiter ‘CLS’ and positional tokens ‘SEP’ but they did not form the “Computer-Pencil” connection.

Original: Furthermore, in Fig.(2b), there are many “Book-Pencil” connections and “Book-Computer” connections that exist in other layers but they all fail to capture the “Computer-Pencil” connection.
Revised: Furthermore, in Fig.(2b), there are many “Book-Pencil” connections and “Book-Computer” connections that exist in other layers but they all fail to capture the “Computer-Pencil” connection.

Original: The fifth head in the third layer is the miracle head that can establish the strong linguistic connection between “computer” and “Pencil” and it forms lexical self-attention patterns across different sentences.
Revised: The fifth head in the third layer is the miracle head that can establish the strong linguistic connection between “computer” and “Pencil” and it forms lexical self-attention patterns across different sentences.

Original: We call this an exact example of the High-order Statistics , and it is formulated on input pairs which satisfy the following definition:
Revised: We call this an exact example of the High-order Statistics , and it is formulated on input pairs which satisfy the following definition:

Original: Definition 3.1.
Revised: Definition 3.1.

Original: ( δ, p, W Q , W K )- similar For vector inputs x i , x j and x k , there exists a single head self-attention with query transformations W Q and key transformations W K , and the selfattention computation forms the feature map score as S ij =
Revised: ( δ, p, W Q , W K )- similar For vector inputs x i , x j and x k , there exists a single head self-attention with query transformations W Q and key transformations W K , and the selfattention computation forms the feature map score as S ij =

Original: ( x i W Q )( x j W K )
Revised: ( x i W Q )( x j W K )

Original: ⊤ such that, under the δ -significant threshold p , if S ik ≥ p and S jk ≥ p, then Pr[ S ij ≥ p ] > δ otherwise, then Pr[ S ij ≥ p ] ≤ 1 − δ .
Revised: ⊤ such that, under the δ -significant threshold p , if S ik ≥ p and S jk ≥ p, then Pr[ S ij ≥ p ] > δ otherwise, then Pr[ S ij ≥ p ] ≤ 1 − δ .

Original: Thus the inputs x i and x j are called ( δ, p, W Q , W K )-similar w.r.t the input x k .
Revised: Thus the inputs x i and x j are called ( δ, p, W Q , W K )-similar w.r.t the input x k .

Original: In Fig.(2c), we relax the restriction on significant connections and select the major connections covering 70.5% local connections among the three words.
Revised: In Fig.(2c), we relax the restriction on significant connections and select the major connections covering 70.5% local connections among the three words.

Original: As a natural result of decreasing the significance δ and threshold p in Definition 3.1, it emerges more high-order self-attentions, i.e., connections between “Computer” and “Pencil”.
Revised: As a natural result of decreasing the significance δ and threshold p in Definition 3.1, it emerges more high-order self-attentions, i.e., connections between “Computer” and “Pencil”.

Original: This implies that weak high-order self-attentions exist in the canonical self-attention but these connections are too delicate to contribute to selfattention’s layer out.
Revised: This implies that weak high-order self-attentions exist in the canonical self-attention but these connections are too delicate to contribute to selfattention’s layer out.

Original: 3.2 Calculate the Jump Self-attention
Revised: 3.2 Calculate the Jump Self-attention

Original: The Jump Self-attention (JAT) is based on the vanilla Transformer architecture to be consistent with the latest improvement on it.
Revised: The Jump Self-attention (JAT) is based on the vanilla Transformer architecture to be consistent with the latest improvement on it.

Original: Inspired by the pieces moving in Fig.(1), we build a transaction graph on the self-attention feature map.
Revised: Inspired by the pieces moving in Fig.(1), we build a transaction graph on the self-attention feature map.

Original: Then JAT utilizes the GCN operator to allow “jump” connections to traverse through different queries and key pairs, and the weak connections under the similar condition are enhanced to help enforces the desired triadic closure property.
Revised: Then JAT utilizes the GCN operator to allow “jump” connections to traverse through different queries and key pairs, and the weak connections under the similar condition are enhanced to help enforces the desired triadic closure property.

Original: Let the attention score of single-headed self-attention in Eq.(1) be S =
Revised: Let the attention score of single-headed self-attention in Eq.(1) be S =

Original: QK ⊤ , we notice that the softmax( · ) function on S / √ d is row-wise, which makes the self-attention layer output the correct rerepresentation of value V .
Revised: QK ⊤ , we notice that the softmax( · ) function on S / √ d is row-wise, which makes the self-attention layer output the correct rerepresentation of value V .

Original: To understand the underlying meaning of high-order statistics in attention computation, we can take a column view of S in Fig.(3a).
Revised: To understand the underlying meaning of high-order statistics in attention computation, we can take a column view of S in Fig.(3a).

Original: Different from the row view where each row reflects one query’s scores on all keys, in the column view, each key column represents a similarity measurement of all queries against itself.
Revised: Different from the row view where each row reflects one query’s scores on all keys, in the column view, each key column represents a similarity measurement of all queries against itself.

Original: Consider an intuitive example from the previous case study, we draw the attention feature map with five queries and keys in Fig.(3b).
Revised: Consider an intuitive example from the previous case study, we draw the attention feature map with five queries and keys in Fig.(3b).

Original: We firstly present the dot-product self-attention scores on the left side.
Revised: We firstly present the dot-product self-attention scores on the left side.

Original: The self-attention scores S 31 = Q 3 ( K 1 ) ⊤ and S 

Revised: The self-attention scores S 31 = Q 3 ( K 1 ) ⊤ and S

Original: =
Revised: =

Original: Q 4 ( K 1 ) ⊤ highlight in the column vector of the first key, as we can see the crosshatch under the column view (the red shade), which indicates that inputs x 3 and x 4 attend to the input x 1 from the K 1 ’s observation.
Revised: Q 4 ( K 1 ) ⊤ highlight in the column vector of the first key, as we can see the crosshatch under the column view (the red shade), which indicates that inputs x 3 and x 4 attend to the input x 1 from the K 1 ’s observation.

Original: The connection between x
Revised: The connection between x

Original: [Figure]
Revised: [Figure]

Original: and x 4 are weak (the blank square).
Revised: and x 4 are weak (the blank square).

Original: Based on the high-order statistics, we have a high probability δ that the inputs x 3 “ Pencil ” and x 4 “ Computer ” becomes ( δ, p, W Q , W K )-similar w.r.t the input x 1 “ Book ”.
Revised: Based on the high-order statistics, we have a high probability δ that the inputs x 3 “ Pencil ” and x 4 “ Computer ” becomes ( δ, p, W Q , W K )-similar w.r.t the input x 1 “ Book ”.

Original: But the dot-product self-attention can not capture this kind of similarity for x 3 and x 4 .
Revised: But the dot-product self-attention can not capture this kind of similarity for x 3 and x 4 .

Original: If we could modify the linear transformation W Q and W K such that the red marked square attains larger attention scores while others maintains, and the attention score between x 3 and x 4 grows larger enough than the significance threshold p .
Revised: If we could modify the linear transformation W Q and W K such that the red marked square attains larger attention scores while others maintains, and the attention score between x 3 and x 4 grows larger enough than the significance threshold p .

Original: We call this attention adjustment a successful “jump”, and it could happen for the individual key.
Revised: We call this attention adjustment a successful “jump”, and it could happen for the individual key.

Original: We assume that there exists a “miracle head”, whose self-attention feature map contains both dotproduct similarity and high-order statistics.
Revised: We assume that there exists a “miracle head”, whose self-attention feature map contains both dotproduct similarity and high-order statistics.

Original: Instead of modifying the linear transformation W Q and W K , we can direct “correct” the self-attention feature map by running thorough all triplets and adjusting the scores by specific keys.
Revised: Instead of modifying the linear transformation W Q and W K , we can direct “correct” the self-attention feature map by running thorough all triplets and adjusting the scores by specific keys.

Original: However, this direct operation will introduce an extra computation of attention combination
Revised: However, this direct operation will introduce an extra computation of attention combination

Original: O
Revised: O

Original: ( L 3 ) w.r.t the input length L .
Revised: ( L 3 ) w.r.t the input length L .

Original: And it is hard to acquire a back-propagation friendly implementation.
Revised: And it is hard to acquire a back-propagation friendly implementation.

Original: To formulate the modification of feature map, we relax the objective of finding the miracle heads to enhance the weak “jump” connection in canonical self-attention computation.
Revised: To formulate the modification of feature map, we relax the objective of finding the miracle heads to enhance the weak “jump” connection in canonical self-attention computation.

Original: Under the row view, we build a graph G =
Revised: Under the row view, we build a graph G =

Original: ( V, E ) from the self-attention scores, where the queries compose the node set
Revised: ( V, E ) from the self-attention scores, where the queries compose the node set

Original: V with the row vector of S as node features and each edge is defined by a adjacency matrix A .
Revised: V with the row vector of S as node features and each edge is defined by a adjacency matrix A .

Original: Then, we introduce A ’s computation.
Revised: Then, we introduce A ’s computation.

Original: Subscribing the matrix with i ⊤ returns the i -th column as a vector.
Revised: Subscribing the matrix with i ⊤ returns the i -th column as a vector.

Original: We select the j -th key as the current column, then the sub-adjacency matrix under the K j is defined:
Revised: We select the j -th key as the current column, then the sub-adjacency matrix under the K j is defined:

Original: [Equation]
Revised: [Equation]

Original: The operation ⊗ denotes the tensor product and we use the similar d -dimension normalization in the canonical self-attention [18].
Revised: The operation ⊗ denotes the tensor product and we use the similar d -dimension normalization in the canonical self-attention [18].

Original: U j is a symmetric matrix containing input pairs with ( δ, p, W Q , W K )- similarity, like x 3 and x 4 in the previous example.
Revised: U j is a symmetric matrix containing input pairs with ( δ, p, W Q , W K )- similarity, like x 3 and x 4 in the previous example.

Original: [ · ] ρ is an indicator function for each matrix item that returns 1 if the item is greater than the ρ , otherwise return 0 .
Revised: [ · ] ρ is an indicator function for each matrix item that returns 1 if the item is greater than the ρ , otherwise return 0 .

Original: Here, we break the “both” requirements in Definition 3.1.
Revised: Here, we break the “both” requirements in Definition 3.1.

Original: Thus, the threshold ρ becomes a hyperparameter reflecting the magnitude of δ and p .
Revised: Thus, the threshold ρ becomes a hyperparameter reflecting the magnitude of δ and p .

Original: We go through all keys, the j -th sub-adjacency matrix reflects the similarity between different queries w.r.t K j .
Revised: We go through all keys, the j -th sub-adjacency matrix reflects the similarity between different queries w.r.t K j .

Original: Inspired by the layer aggregation [17, 10], the summation of weighted adjacency matrix yield larger gap between isolated values and bulk eigenvalues, which has been widely used in helping the community detection task [14, 4].
Revised: Inspired by the layer aggregation [17, 10], the summation of weighted adjacency matrix yield larger gap between isolated values and bulk eigenvalues, which has been widely used in helping the community detection task [14, 4].

Original: The sub-adjacency matrices are describing the connections between all the queries, i.e, different layers in a multiplex network, we can also leverage the same operation for a better edge probability recovery – similarity detection.
Revised: The sub-adjacency matrices are describing the connections between all the queries, i.e, different layers in a multiplex network, we can also leverage the same operation for a better edge probability recovery – similarity detection.

Original: Then, we can add up all the sub-adjacency matrix into the weighted one:
Revised: Then, we can add up all the sub-adjacency matrix into the weighted one:

Original: [Equation]
Revised: [Equation]

Original: Our next question is to enhance the “jump” connection for individual inputs.
Revised: Our next question is to enhance the “jump” connection for individual inputs.

Original: The Graph Convolutional
Revised: The Graph Convolutional

Original: Neural network (GCN) could perform local aggregate operators over the graph adjacency matrix, which is proposed by Scarselli et al.
Revised: Neural network (GCN) could perform local aggregate operators over the graph adjacency matrix, which is proposed by Scarselli et al.

Original: [15] and is further simplified by Duvenaud et al.
Revised: [15] and is further simplified by Duvenaud et al.

Original: [6] and
Revised: [6] and

Original: Sukhbaatar et al.
Revised: Sukhbaatar et al.

Original: We employ it as a specific architecture for encouraging the attended items to become similar in both queries and keys.
Revised: We employ it as a specific architecture for encouraging the attended items to become similar in both queries and keys.

Original: Then we have:
Revised: Then we have:

Original: [Equation]
Revised: [Equation]

Original: where ˆ A =
Revised: where ˆ A =

Original: ˜ D − 12 ˜
Revised: ˜ D − 12 ˜

Original: A ˜ D − 12 , ˜
Revised: A ˜ D − 12 , ˜

Original: A =
Revised: A =

Original: A + I and ˜ D is a diagonal matrix from the column-sum of A , where every node can only attend to its nearest neighborhood with the learnable graph filter W .
Revised: A + I and ˜ D is a diagonal matrix from the column-sum of A , where every node can only attend to its nearest neighborhood with the learnable graph filter W .

Original: The
Revised: The

Original: GCN will make the selected queries/keys propagate its features, which causes a higher score when encountering each other in self-attention computation.
Revised: GCN will make the selected queries/keys propagate its features, which causes a higher score when encountering each other in self-attention computation.

Original: In other words, GCN encourages high-order self-attention’s propagation through individual keys’ similarity measurement.
Revised: In other words, GCN encourages high-order self-attention’s propagation through individual keys’ similarity measurement.

Original: Thanks to this property, the self-attention score will be updated as:
Revised: Thanks to this property, the self-attention score will be updated as:

Original: [Equation]
Revised: [Equation]

Original: If we aggregate the graph filter W into the random projection W as a new projection W ′ , the score S ′ becomes ˆ A ( XW ′ Q )( XW ′ K )
Revised: If we aggregate the graph filter W into the random projection W as a new projection W ′ , the score S ′ becomes ˆ A ( XW ′ Q )( XW ′ K )

Original: ⊤ ˆ A ⊤ .
Revised: ⊤ ˆ A ⊤ .

Original: This could be consider as a way to modify the original learnable matrix W Q into the new projection W ′ Q .
Revised: This could be consider as a way to modify the original learnable matrix W Q into the new projection W ′ Q .

Original: Note that the proposed method maintains the same parameter scale.
Revised: Note that the proposed method maintains the same parameter scale.

Original: Immediately, we could derive the self-attention score after applying the GCN operator on both queries and keys, namely the “jump” operation:
Revised: Immediately, we could derive the self-attention score after applying the GCN operator on both queries and keys, namely the “jump” operation:

Original: [Equation]
Revised: [Equation]

Original: which forms an approximated graph operator repeating propagation over the graph to enforce the desired triadic closure, like SGC [22]
Revised: which forms an approximated graph operator repeating propagation over the graph to enforce the desired triadic closure, like SGC [22]

Original: The original GCN activation operation is omitted to avoid interfering with later self-attention calculations, which is merged into the following Feed-Forward
Revised: The original GCN activation operation is omitted to avoid interfering with later self-attention calculations, which is merged into the following Feed-Forward

Original: Net.
Revised: Net.

Original: In this way, we are able to propose the Jump Self-attention:
Revised: In this way, we are able to propose the Jump Self-attention:

Original: [Equation]
Revised: [Equation]

Original: To better take advantage of the pre-trained model, we suggest to split H heads in canonical selfattention into two groups of H c heads and H j heads.
Revised: To better take advantage of the pre-trained model, we suggest to split H heads in canonical selfattention into two groups of H c heads and H j heads.

Original: The H c heads perform the canonical selfattention operation and the H j heads perform jump self-attention, then we concatenate their results as the complete layer output.
Revised: The H c heads perform the canonical selfattention operation and the H j heads perform jump self-attention, then we concatenate their results as the complete layer output.

Original: The use of high-order self-attention enables acquiring high-order dependencies over the superficial ones.
Revised: The use of high-order self-attention enables acquiring high-order dependencies over the superficial ones.

Original: Overview of JAT’s whole process is given in Appendix.
Revised: Overview of JAT’s whole process is given in Appendix.

Original: The
Revised: The

Original: JAT and canonical self-attention can be used interchangeably in Transformer models.
Revised: JAT and canonical self-attention can be used interchangeably in Transformer models.

Original: 3.2.1 Masking tricks
Revised: 3.2.1 Masking tricks

Original: The mask of self-attention [18] plays an important role in avoiding the auto-regressive during the inference phase.
Revised: The mask of self-attention [18] plays an important role in avoiding the auto-regressive during the inference phase.

Original: The GCN operations may cause leftward information to leak between queries.
Revised: The GCN operations may cause leftward information to leak between queries.

Original: If
Revised: If

Original: JAT is applied in the decoder, we have to prevent this in building the adjacency matrix, where we reduce the edges from afterward queries to previous queries.
Revised: JAT is applied in the decoder, we have to prevent this in building the adjacency matrix, where we reduce the edges from afterward queries to previous queries.

Original: We implement this inside of JAT by forcing the upper triangular of adjacency matrix to 0 in Eq.(3).
Revised: We implement this inside of JAT by forcing the upper triangular of adjacency matrix to 0 in Eq.(3).

Original: The canonical self-attention’s masking operation on the feature map is preserved after applying GCN in Eq.(5) such that we can eschew the information leakage in softmax( · ) function.
Revised: The canonical self-attention’s masking operation on the feature map is preserved after applying GCN in Eq.(5) such that we can eschew the information leakage in softmax( · ) function.

Original: 3.2.2 Efficient variants
Revised: 3.2.2 Efficient variants

Original: The computation complexity of JAT remains O ( L 3 ) w.r.t the input length L , because we have to traverse all keys in the adjacency matrix aggregation of Eq.(3).
Revised: The computation complexity of JAT remains O ( L 3 ) w.r.t the input length L , because we have to traverse all keys in the adjacency matrix aggregation of Eq.(3).

Original: We previously designed a sparsity measurement [24] to find the most significant queries in self-attention, which motivates us to employ it in selecting the primary keys:
Revised: We previously designed a sparsity measurement [25] to find the most significant queries in self-attention, which motivates us to employ it in selecting the primary keys:

Original: [Equation]
Revised: [Equation]

Original: We follow the sampling strategy and choose Top- u Keys.
Revised: We follow the sampling strategy and choose Top- u Keys.

Original: Thus, the original self-attention S reduces to ¯ S ∈ R L × u
Revised: Thus, the original self-attention S reduces to ¯ S ∈ R L × u

Original: and we acquire the efficient variant JAT † in the O
Revised: and we acquire the efficient variant JAT † in the O

Original: ( L 2 ) .
Revised: ( L 2 ) .

Original: Recalling that the JAT shares the same parameter scale, and the JAT † holds the same computation complexity when compared with the canonical self-attention.
Revised: Recalling that the JAT shares the same parameter scale, and the JAT † holds the same computation complexity when compared with the canonical self-attention.

Original: [Table]
Revised: [Table]

Original: 3.2.3 Higher-order variants
Revised: 3.2.3 Higher-order variants

Original: Without loss of generality
Revised: Without loss of generality

Original: [Table]
Revised: [Table]

Original: we can extend the similar in Definition 3.1 to a higher order case
Revised: we can extend the similar in Definition 3.1 to a higher order case

Original: [Table]
Revised: [Table]

Original: where the similarity depends on more than one input.
Revised: where the similarity depends on more than one input.

Original: Instead of performing analogical derivations, we can introduce the Mix-hop [1] GCN to capture the multi-hop propagation on the graph of queries.
Revised: Instead of performing analogical derivations, we can introduce the Mix-hop [1] GCN to capture the multi-hop propagation on the graph of queries.

Original: If we refer the canonical self-attention [18] to the first order self-attention and our JAT to the second order one, then we can define the j -th order self-attention as:
Revised: If we refer the canonical self-attention [18] to the first order self-attention and our JAT to the second order one, then we can define the j -th order self-attention as:

Original: [Equation]
Revised: [Equation]

Original: where ˆ A j − 1 represents the adjacency matrix ˆ A multiplied by itself ( j − 1) times.
Revised: where ˆ A j − 1 represents the adjacency matrix ˆ A multiplied by itself ( j − 1) times.

Original: It’s trivial to substitute the Φ j ( S ) in Eq.(7) for higher-order self-attention variants A 2 , A 3 and etc.
Revised: It’s trivial to substitute the Φ j ( S ) in Eq.(7) for higher-order self-attention variants A 2 , A 3 and etc.

Original: We can split the heads into more groups when concatenating multiple self-attentions with different orders.
Revised: We can split the heads into more groups when concatenating multiple self-attentions with different orders.

Original:  Experiments
Revised:  Experiments

Original: In this section, we empirically demonstrate JAT’s effectiveness on the pre-trained BERT base model with 12 layers, and perform BERT-JAT fine-tuning and evaluation on ten NLP tasks.
Revised: In this section, we empirically demonstrate JAT’s effectiveness on the pre-trained BERT base model with 12 layers, and perform BERT-JAT fine-tuning and evaluation on ten NLP tasks.

Original: 4.1 Setup
Revised: 4.1 Setup

Original: (a) GLUE.
Revised: (a) GLUE.

Original: We conduct JAT’s experiments of the language understanding and generalization capabilities on the General Language Understanding Evaluation (GLUE) benchmark [21], a collection of diverse natural language understanding tasks.
Revised: We conduct JAT’s experiments of the language understanding and generalization capabilities on the General Language Understanding Evaluation (GLUE) benchmark [21], a collection of diverse natural language understanding tasks.

Original: We also perform additional experiments on SuperGLUE benchmark [20].
Revised: We also perform additional experiments on SuperGLUE benchmark [20].

Original: Settings: We use a batch size of 32 and fine-tune for 5 epochs over the data for nine GLUE tasks.
Revised: Settings: We use a batch size of 32 and fine-tune for 5 epochs over the data for nine GLUE tasks.

Original: The threshold ρ is selected from { 4 . 0 , 4 . 5 , . . . , 10 . 0 } .
Revised: The threshold ρ is selected from { 4 . 0 , 4 . 5 , . . . , 10 . 0 } .

Original: The other settings follow the recommendation of the original paper.
Revised: The other settings follow the recommendation of the original paper.

Original: Since the proposed JAT can be used interchangeably with canonical self-attention, we perform a grid search on the layer replacement.
Revised: Since the proposed JAT can be used interchangeably with canonical self-attention, we perform a grid search on the layer replacement.

Original: There are two sets of layer deployment, where the first combination is chosen from {Layer 1 − 4 ,
Revised: There are two sets of layer deployment, where the first combination is chosen from {Layer 1 − 4 ,

Original: Layer 5 − 8 , Layer 9 − 12 } and the alternative is {Layer 1 − 6 , Layer 7 − 12 }.
Revised: Layer 5 − 8 , Layer 9 − 12 } and the alternative is {Layer 1 − 6 , Layer 7 − 12 }.

Original: Another important selection is the multi-heads grouping, we employ the “side-by-side” strategy as replacing the heads {2, 4, 6, 8, 10} with JAT.
Revised: Another important selection is the multi-heads grouping, we employ the “side-by-side” strategy as replacing the heads {2, 4, 6, 8, 10} with JAT.

Original: And we do not use any ensembling strategy or multi-tasking scheme in this fine-tuning.
Revised: And we do not use any ensembling strategy or multi-tasking scheme in this fine-tuning.

Original: The evaluation is performed on the Dev set.
Revised: The evaluation is performed on the Dev set.

Original: Metric: We use three different evaluation metrics on the 9 tasks.
Revised: Metric: We use three different evaluation metrics on the 9 tasks.

Original: Matthews correlation coefficient ( MCC =
Revised: Matthews correlation coefficient ( MCC =

Original: [Equation]
Revised: [Equation]

Original: √
Revised: √

Original: [Equation]
Revised: [Equation]

Original: ( TP + FN )(
Revised: ( TP + FN )(

Original: [Equation]
Revised: [Equation]

Original: + FP )( TN + FN ) ) is used for CoLA, Pearson ( PCC =
Revised: + FP )( TN + FN ) ) is used for CoLA, Pearson ( PCC =

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: ) is used for STS-B, and accuracy is used for others.
Revised: ) is used for STS-B, and accuracy is used for others.

Original: Platform: Intel Xeon 3.2GHz + The Nvidia V100 GPU (32 GB) X 4.
Revised: Platform: Intel Xeon 3.2GHz + The Nvidia V100 GPU (32 GB) X 4.

Original: The code is available at https://github.com/zhouhaoyi/JAT2022 .
Revised: The code is available at https://github.com/zhouhaoyi/JAT2022 .

Original: (b) SQuAD.
Revised: (b) SQuAD.

Original: We also evaluate on the SQuAD v1.1/v2.0 [12, 13].
Revised: We also evaluate on the SQuAD v1.1/v2.0 [12, 13].

Original: The settings is in Appendix D
Revised: The settings is in Appendix D

Original: 4.2 Main results
Revised: 4.2 Main results

Original: The result of the GLUE experiment is summarized in Table 1.
Revised: The result of the GLUE experiment is summarized in Table 1.

Original: We compare ELMo [8] the baseline of GLUE, BERT base [5], and RoBERTa [9], the pre-trained model where we build JAT varints.
Revised: We compare ELMo [8] the baseline of GLUE, BERT base [5], and RoBERTa [9], the pre-trained model where we build JAT varints.

Original: Our model BERT/RoBERTa-JAT outperform base models on eight tasks without introducing more
Revised: Our model BERT/RoBERTa-JAT outperform base models on eight tasks without introducing more

Original: [Table] [Table]
Revised: [Table] [Table]

Original: parameters, which shows that JAT can capture higher-order dependencies than canonical self-attention.
Revised: parameters, which shows that JAT can capture higher-order dependencies than canonical self-attention.

Original: More specifically
Revised: More specifically

Original: [Table]
Revised: [Table]

Original: BERT-JAT achieve 9.4% score rising on the CoLA dataset and 7.1% on the WNLI dataset as a remarkable improvement, which shows that JAT shows the better alignment ability on a smaller dataset.
Revised: BERT-JAT achieve 9.4% score rising on the CoLA dataset and 7.1% on the WNLI dataset as a remarkable improvement, which shows that JAT shows the better alignment ability on a smaller dataset.

Original: When it applies to RoBERTa, the average gain is larger and the most significant rising 4.9% happens on the CoLA dataset, which demonstrates a stable performance gain on SOTA models.
Revised: When it applies to RoBERTa, the average gain is larger and the most significant rising 4.9% happens on the CoLA dataset, which demonstrates a stable performance gain on SOTA models.

Original: We have also show that JAT beats the competitor
Revised: We have also show that JAT beats the competitor

Original: A 3 [23] generally.
Revised: A 3 [24] generally.

Original: The A 3 introduces the triplet attention measuring dissimilarity to diversify the self-attention, while our proposed JAT tried to enhance similarity by building connections between unrelated pairs from a different perspective.
Revised: The A 3 introduces the triplet attention measuring dissimilarity to diversify the self-attention, while our proposed JAT tried to enhance similarity by building connections between unrelated pairs from a different perspective.

Original: These two works are orthogonal and can be performed together to improve self-attention mechanism.
Revised: These two works are orthogonal and can be performed together to improve self-attention mechanism.

Original: Meanwhile, the efficient variant JAT † also shows competitive performance than competitor A 3 .
Revised: Meanwhile, the efficient variant JAT † also shows competitive performance than competitor A 3 .

Original: Results of evaluating JAT on the more challenging SuperGLUE benchmark are presented in Table 2, we have selected RoBERTa as a more efficient basic model.
Revised: Results of evaluating JAT on the more challenging SuperGLUE benchmark are presented in Table 2, we have selected RoBERTa as a more efficient basic model.

Original: With more extended tokens on complicated tasks, the JAT fails on three tasks for OOM, but it still acquires better results than vanilla
Revised: With more extended tokens on complicated tasks, the JAT fails on three tasks for OOM, but it still acquires better results than vanilla

Original: RoBERTa on others.
Revised: RoBERTa on others.

Original: The efficient variant RoBERTa-JAT † achieves the best average scores.
Revised: The efficient variant RoBERTa-JAT † achieves the best average scores.

Original: The result of the SQuAD experiment is summarized in Table 3.
Revised: The result of the SQuAD experiment is summarized in Table 3.

Original: Our model BERT-JAT achieves better performance than all baselines in both EM and F1 scores, which is aligned with the above results.
Revised: Our model BERT-JAT achieves better performance than all baselines in both EM and F1 scores, which is aligned with the above results.

Original: 4.3 Ablation Study 4.3.1 Effect of JAT’s Layer Deployment
Revised: 4.3 Ablation Study 4.3.1 Effect of JAT’s Layer Deployment

Original: In this experiment, we explore the JAT layer’s effect on acquiring the high-order connections.
Revised: In this experiment, we explore the JAT layer’s effect on acquiring the high-order connections.

Original: We set the layers at a different position in RoBERTa base model to use JAT heads, where the RoBERTaJAT l [ i,j ] represents that the layers from i to j in RoBERTa base model have four JAT heads.
Revised: We set the layers at a different position in RoBERTa base model to use JAT heads, where the RoBERTaJAT l [ i,j ] represents that the layers from i to j in RoBERTa base model have four JAT heads.

Original: We report the scores on the selected GLUE tasks, including CoLA, MRPC, RTE, and STS-B in Table 4.
Revised: We report the scores on the selected GLUE tasks, including CoLA, MRPC, RTE, and STS-B in Table 4.

Original: We can conclude that using JAT heads in lower or middle layers leads to better scores.
Revised: We can conclude that using JAT heads in lower or middle layers leads to better scores.

Original: This is consistent with the findings in Fig.(2), middle layers possess more weak high-order connections to be easily enhanced.
Revised: This is consistent with the findings in Fig.(2), middle layers possess more weak high-order connections to be easily enhanced.

Original: Applying JAT at lower layers can effectively enhance the diversity of lower layers and improve the high-order connections to force the desired triadic closure.
Revised: Applying JAT at lower layers can effectively enhance the diversity of lower layers and improve the high-order connections to force the desired triadic closure.

Original: 4.3.2 Effect of JAT’s Head Grouping
Revised: 4.3.2 Effect of JAT’s Head Grouping

Original: To evaluate the effect of JAT’s head grouping (different splits) on general language understanding task accuracy, we trained many RoBERTa-JAT models with different JAT head numbers in layer 1-6.
Revised: To evaluate the effect of JAT’s head grouping (different splits) on general language understanding task accuracy, we trained many RoBERTa-JAT models with different JAT head numbers in layer 1-6.

Original: Results on the selected GLUE tasks are shown in Table 5.
Revised: Results on the selected GLUE tasks are shown in Table 5.

Original: With the number of JAT head increasing, we can find that almost all the tasks’ scores first increase then decrease.
Revised: With the number of JAT head increasing, we can find that almost all the tasks’ scores first increase then decrease.

Original: The scores increase because adding JAT heads enhance the high-order global dependencies in the vanilla RoBERTa model.
Revised: The scores increase because adding JAT heads enhance the high-order global dependencies in the vanilla RoBERTa model.

Original: They decrease when the number of heads is greater than 6, which probably is caused by too strong highorder inductive bias suppressing the canonical self-attention.
Revised: They decrease when the number of heads is greater than 6, which probably is caused by too strong highorder inductive bias suppressing the canonical self-attention.

Original: But the CoLA dataset seems fitting the
Revised: But the CoLA dataset seems fitting the

Original: JAT assumption very well, which may explain why it reaches the highest performance gain in Table 1.
Revised: JAT assumption very well, which may explain why it reaches the highest performance gain in Table 1.

Original: 4.3.3 Effect of Higher Order in JAT
Revised: 4.3.3 Effect of Higher Order in JAT

Original: We perform this experiment to evaluate the effects of higher-order variants.
Revised: We perform this experiment to evaluate the effects of higher-order variants.

Original: Following the similar settings as the previous ablation study, where the RoBERTa-JAT oi stands for the i -th order selfattention defined in Eq.(9).
Revised: Following the similar settings as the previous ablation study, where the RoBERTa-JAT oi stands for the i -th order selfattention defined in Eq.(9).

Original: We show the selected GLUE results in Table 6.
Revised: We show the selected GLUE results in Table 6.

Original: And it reveals that for most natural language understanding tasks, the second-order self-attention is sufficient to discover
Revised: And it reveals that for most natural language understanding tasks, the second-order self-attention is sufficient to discover

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: more higher-order connections and improve the expression ability of the model, which leads to better results.
Revised: more higher-order connections and improve the expression ability of the model, which leads to better results.

Original: For example, using the second order JAT on RTE dataset significantly improves the scores.
Revised: For example, using the second order JAT on RTE dataset significantly improves the scores.

Original: 4.3.4 Effect of Layer Stacking Degradation in JAT
Revised: 4.3.4 Effect of Layer Stacking Degradation in JAT

Original: Similar to the evaluation in A 3 [23], we conduct layer stacking degradation experiments on both
Revised: Similar to the evaluation in A 3 [24], we conduct layer stacking degradation experiments on both

Original: CoLA and RTE downstream tasks.
Revised: CoLA and RTE downstream tasks.

Original: Firstly, we reduce the simpler BERT base model by one layer at a time, which makes the overall layers degrade from 12 to 2.
Revised: Firstly, we reduce the simpler BERT base model by one layer at a time, which makes the overall layers degrade from 12 to 2.

Original: Fig.(4) shows that the BERT-JAT consistently outperforms BERT base .
Revised: Fig.(4) shows that the BERT-JAT consistently outperforms BERT base .

Original: In the RTE task, the BERT-JAT even achieves the highest accuracy with eight layers.
Revised: In the RTE task, the BERT-JAT even achieves the highest accuracy with eight layers.

Original: Secondly, we have extended the BERT base model from 12 layers tolayers, which matches the comparable extra computation brought by the JAT mechanism (refer to
Revised: Secondly, we have extended the BERT base model from 12 layers tolayers, which matches the comparable extra computation brought by the JAT mechanism (refer to

Original: Sec 4.6).
Revised: Sec 4.6).

Original: It reveals that the deeper baseline lags behind the JAT-enhanced ones.
Revised: It reveals that the deeper baseline lags behind the JAT-enhanced ones.

Original: 4.4 Visualization
Revised: 4.4 Visualization

Original: We conduct a visualization in Fig.(5) to show the self-attention feature map changing in JAT.
Revised: We conduct a visualization in Fig.(5) to show the self-attention feature map changing in JAT.

Original: To avoid repetition, we construct a new sentence.
Revised: To avoid repetition, we construct a new sentence.

Original: The line between words represents their attention scores: the darker, the larger.
Revised: The line between words represents their attention scores: the darker, the larger.

Original: The Fig.(5a) shows heads of layer 6 in BERT base , and it indicates that the “Kate” has stronger connections to the words in the second sentence.
Revised: The Fig.(5a) shows heads of layer 6 in BERT base , and it indicates that the “Kate” has stronger connections to the words in the second sentence.

Original: The Fig.(5b) shows heads of the same layer in BERT-JAT, the “Kate” attends to the first sentence’s “hit” through a “jump” connection with the blue head.
Revised: The Fig.(5b) shows heads of the same layer in BERT-JAT, the “Kate” attends to the first sentence’s “hit” through a “jump” connection with the blue head.

Original: This example shows that ‘Kate’ never gets a strong connection with the action ‘hit’ in vanilla self-attention, but JAT provides this conditional connection (may through ‘bite’), which indicates the two actions may be sequentially related.
Revised: This example shows that ‘Kate’ never gets a strong connection with the action ‘hit’ in vanilla self-attention, but JAT provides this conditional connection (may through ‘bite’), which indicates the two actions may be sequentially related.

Original: This JAT-inspired attention change also applies to the ‘SEP’ token, which implies that ‘Kate’ are the last few words.
Revised: This JAT-inspired attention change also applies to the ‘SEP’ token, which implies that ‘Kate’ are the last few words.

Original: Although this result differs from the previous examples on the noun exchanging, it reveals that the Jump Self-attention can enhance connections with various lexical patterns.
Revised: Although this result differs from the previous examples on the noun exchanging, it reveals that the Jump Self-attention can enhance connections with various lexical patterns.

Original: 4.5 Case Study
Revised: 4.5 Case Study

Original: We conduct an experiment on Named-Entity Recognition (NER) task using the standard CoNLLdataset, which concentrates on four types of named entities: persons, locations, organizations and names of miscellaneous entities that do not belong to the previous three groups.
Revised: We conduct an experiment on Named-Entity Recognition (NER) task using the standard CoNLLdataset, which concentrates on four types of named entities: persons, locations, organizations and names of miscellaneous entities that do not belong to the previous three groups.

Original: The settings are included in Appendix E. We select the major connections in the attention feature map where their scores are greater than (mean + std).
Revised: The settings are included in Appendix E. We select the major connections in the attention feature map where their scores are greater than (mean + std).

Original: We use ‘CROSS’ to represent connections of tokens with different named-entities.
Revised: We use ‘CROSS’ to represent connections of tokens with different named-entities.

Original: The visualization is showed in Fig.(7).
Revised: The visualization is showed in Fig.(7).

Original: From Fig.(7a) we find there are not many ‘CROSS’ connections in the attention of BERT base , and most of the connections are between tokens belong to person entity.
Revised: From Fig.(7a) we find there are not many ‘CROSS’ connections in the attention of BERT base , and most of the connections are between tokens belong to person entity.

Original: From Fig.(7b), we observe JAT significantly increases the number of connections between tokens with different named-entities at the layers {1,2,3,8,9,10} having JAT heads, which enhances the model’s ability to discover high-order information and makes the model’s expression more diverse.
Revised: From Fig.(7b), we observe JAT significantly increases the number of connections between tokens with different named-entities at the layers {1,2,3,8,9,10} having JAT heads, which enhances the model’s ability to discover high-order information and makes the model’s expression more diverse.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: our model BERT-JAT improve the connection between tokens with different named-entities (the ‘CROSS’ connections in orange).
Revised: our model BERT-JAT improve the connection between tokens with different named-entities (the ‘CROSS’ connections in orange).

Original: 4.6 Run-time Comparison
Revised: 4.6 Run-time Comparison

Original: With the same setting, we evaluate the iterative efficiency of JAT on eight downstream task, which has no learnable parameters.
Revised: With the same setting, we evaluate the iterative efficiency of JAT on eight downstream task, which has no learnable parameters.

Original: The Fig.(6) shows that JAT mostly brings about 30% extra computation than the vanilla Transformer, where RTE achieves 75% at most.
Revised: The Fig.(6) shows that JAT mostly brings about 30% extra computation than the vanilla Transformer, where RTE achieves 75% at most.

Original: This variance is caused by the different input lengths L when calculating the adjacency matrix in Eq.(3).
Revised: This variance is caused by the different input lengths L when calculating the adjacency matrix in Eq.(3).

Original: It could be further improved by utilizing the Keys’ sparsity [24], then JAT † achieves no more than 15% extra computation time.
Revised: It could be further improved by utilizing the Keys’ sparsity [25], then JAT † achieves no more than 15% extra computation time.

Original:  Conclusions
Revised:  Conclusions

Original: In this paper, we discovered the existence of Jump Self-attention (JAT) in the Transformers’ dotproduct computation.
Revised: In this paper, we discovered the existence of Jump Self-attention (JAT) in the Transformers’ dotproduct computation.

Original: Drawing inspirations from the pieces moving of English Draughts, we define
Revised: Drawing inspirations from the pieces moving of English Draughts, we define

Original: JAT from the canonical self-attention and propose a tangible framework to formalize JAT into the graph learning problem and leverage it to encourage JAT’s propagation in attention feature map.
Revised: JAT from the canonical self-attention and propose a tangible framework to formalize JAT into the graph learning problem and leverage it to encourage JAT’s propagation in attention feature map.

Original: We further develop the higher-order variants under multi-hop assumption to increase the generality.
Revised: We further develop the higher-order variants under multi-hop assumption to increase the generality.

Original: The proposed architecture is compatible with the pre-trained models.
Revised: The proposed architecture is compatible with the pre-trained models.

Original: Experiments on two tasks demonstrate JAT’s superior performance on capturing high-order dependencies.
Revised: Experiments on two tasks demonstrate JAT’s superior performance on capturing high-order dependencies.

Original: We also want to use
Revised: We also want to use

Original: JAT on MindSpore ( https://www.mindspore.cn/ ), which is a new deep learning computing framework.
Revised: JAT on MindSpore ( https://www.mindspore.cn/ ), which is a new deep learning computing framework.

Original: Acknowledgments
Revised: Acknowledgments

Original: This work was supported by grants from the Natural Science Foundation of China (62202029 and
Revised: This work was supported by grants from the Natural Science Foundation of China (62202029 and

Original: U20B2053).
Revised: U20B2053).

Original: Thanks for the computing infrastructure provided by Beijing Advanced Innovation Center for Big Data and Brain Computing.
Revised: Thanks for the computing infrastructure provided by Beijing Advanced Innovation Center for Big Data and Brain Computing.

Original: This work was also sponsored by CAAI-Huawei MindSpore
Revised: This work was also sponsored by CAAI-Huawei MindSpore

Original: Open Fund.
Revised: Open Fund.

Original: We gratefully acknowledge the support of MindSpore, CANN (Compute Architecture for Neural Networks) and Ascend AI Processor used for this research.
Revised: We gratefully acknowledge the support of MindSpore, CANN (Compute Architecture for Neural Networks) and Ascend AI Processor used for this research.

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] We have include a complete Introduction section.
Revised: [Yes] We have include a complete Introduction section.

Original: Please refer to Sec.1.
Revised: Please refer to Sec.1.

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] We include that in the Limitations / Broad Impact Section.
Revised: [Yes] We include that in the Limitations / Broad Impact Section.

Original: Please refer to Sec.
Revised: Please refer to Sec.

Original: Appendix G (c) Did you discuss any potential negative societal impacts of your work?
Revised: Appendix G. (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] We include that in the Limitations / Broad Impact Section.
Revised: [Yes] We include that in the Limitations / Broad Impact Section.

Original: Please refer to Sec.
Revised: Please refer to Sec.

Original: Appendix G (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: Appendix G. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] We have go through them.
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] We have clarified it in the Sec.3 (b) Did you include complete proofs of all theoretical results?
Revised: [Yes] We have clarified it in the Sec. (b) Did you include complete proofs of all theoretical results?

Original: [N/A] 3.
Revised: [N/A] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] We provide all of them in the supplemental material and URL.
Revised: [Yes] We provide all of them in the supplemental material and URL.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] Please refer to Sec.4.1.
Revised: [Yes] Please refer to Sec.4.1.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [N/A] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: [N/A] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] Please refer to Sec.4.1.
Revised: [Yes] Please refer to Sec.4.1.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re us- ing/curating?
Revised: [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: [Yes] (e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content?
Revised: [Yes] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A] All the dataset are public.
Revised: [N/A] All the dataset are public.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?[N/A] Appendix A Related Work
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: The reading comprehension [32], question and answering [28], sentiment analysis [25], image captioning [31] and other tasks encourage linguistic-meaningful solution [21] with higher-level understanding of the language other than superficial correspondences between inputs and outputs.

Original: The giant model [2] shows the superior performance on various down-stream tasks through the large-scale pre-training.

Original: Recently, CLIP [35, 34],

Original: CogView[30] and DALL · E [38] learn a multi-modal embedding space to estimate the semantic similarity between texts and images, and a handful of prompt techniques [33, 39] explores the remapping of semantic representation.

Original: They reveal the possibility of capturing semantic relations based on the self-attention mechanism.

Original: Others apply the knowledge graph in Transformers [42, 26], but it is resource-hungry and needs human-aid crafting.

Original: The most related work is Triplet Attention [23], it utilizes the diversity to enhance the attention feature map and focus on the attention’s dissimilarity.

Original: However, we aim to measure high-order statistics in attention’s similarity to enhance the dot-product self-attention.

Original: B Mathematical notions

Original: For better understanding, we collected some definition of the matrices in Table 7.

Original: Besides, the miracle heads refer to these most effective heads in capturing the jump self-attention, which is hardly acquired in the standard

Original: Transformer architecture.

Original: If the magnitude of jump connections ρ is fixed, finding the statistics are combination problem.

Original: But the optimization becomes intractable.

Original: [Table]
Revised: [N/A]

Original: C The model architecture

Original: We present the overall architecture of proposed Jump Self-attention model in Fig.(8), especially on the heads splits and connection flows.

Original: To simplify the comparison, we use the same feature map between the canonical self-attention and jump one.

Original: They may be different from individual projects in practice.

Original: The red rectangle refers to the enhanced “jump” connections from JAT.

Original: SplitSplitSplitSplit

Original: [Table]

Original: [Figure]

Original: D Experiments settings of SQuAD

Original: We also conduct JAT’s experiments on the performance of the question-answering task.

Original: The Stanford Question Answering Dataset (SQuAD v1.1/v2.0) [12, 13] is a collection of 150k crowd-sourced question/answer pairs.

Original: The task is to predict the answer text span in the passage given a question and a passage from Wikipedia containing the answer.

Original: Settings: We maintain the same fine-tuning strategy as in the GLUE experiment.

Original: Metric:

Original: Exact match (EM), which is the number of exactly correct answers, and F1 scores, which captures the precision and recall that words chosen as part of the answer are part of the answer.

Original: Platform: Intel Xeon 3.2GHz + The

Original: Nvidia V100 GPU (32 GB) X 4.

Original: E Experiments details on the Case Study E.1 Setup

Original: CoNLL-2003: We conduct a case study experiment on Named-Entity Recognition (NER) task using the standard CoNLL-2003 dataset, which concentrates on four types of named entities: persons, locations, organizations and names of miscellaneous entities that do not belong to the previous three groups.

Original: Settings: We use a batch size of 32 and fine-tune a pre-trained BERT base model and a BERT-JAT model whose layers {1,2,3,8,9,10} have 6 JAT heads for 5 epochs.

Original: The other settings follow past works’ recommendation.

Original: Metric: F1 scores and accuracy are used.

Original: Platform: Intel Xeon 3.2GHz + The Nvidia V100 GPU (32 GB) X 4.

Original: E.2 Results

Original: We preprocess the texts before we feed them to the model, including adding ‘CLS’ and ‘SEP’ special tokens whose named-entity labels are set to -100, and split the words into subwords.

Original: We select the major connections in attention feature map where their scores are greater than (mean + std).

Original: We studied the named-entities of the pair-wise tokens corresponding to those major connections and use ‘MISCs’,‘PERs’,‘ORGs’,‘LOCs’ to represent the pair-wise attention connections of tokens with the same named-entities, which means that both tokens are miscellaneous, persons, locations or organizations.

Original: We use ‘SEP’ to represent connections of two tokens where one of them is different entities.

Original: [Figure]

Original: n t i o n

Original: C o nn

Original: [Figure]

Original: c t i o n

Original: [Figure]

Original: We visualize the results containing the ‘SEP’ connection in Fig.(9).

Original: From Fig.(9a), we can find that the number of ‘SEP’ connection is more than other connections, which means that there are many tokens connected to special tokens ‘CLS’ or ‘SEP’ in self-attention, and these connections have redundant and over-expressed information.

Original: This observation is also consistent with previous research [19, 3].

Original: From Fig.(9b), we can find that JAT reduce the number of ‘SEP’ connection at the layers {1,2,3}, noted that the number of JAT heads in these layers is 6, only half of the total attention heads, which shows that JAT can reduce redundant information in attention and enhance the model’s expression ability.

Original: We further remove the ‘SEP’ connection and visualize the results in Fig.(7).

Original: From Fig.(7a) we can find that there are not many ‘CROSS’ connections in the attention of BERT base , and most of the connections are between tokens that belong to a person’s entity.

Original: From Fig.(7b), we can observe that JAT significantly increase the number of connections between tokens with different named-entities at the layers {1,2,3,8,9,10} having JAT heads, which enhances the model’s ability to discover high-order information and makes the model’s representation more diverse.

Original: E.3 Discussion

Original: We summarized the analysis of high-order statistics in Sec 3.1.

Original: The Fig.(2) is a concrete example.

Original: And the NER case study in Sec 4.5 can be considered as an empirical evaluation of capacity limitation.

Original: The ‘CROSS’ connections represent the high-order connections, which are rare in BERT base while are more in BERT-JAT.

Original: F The performance of training from scratch

Original: Without the fine-tuning framework, where Roberta is a representative SOTA, we want to separately investigate the JAT’s performance when training from scratch.

Original: We have performed experiments at different data scales on

Original: CoLA.

Original: Table 8 shows that the model receives constant gains after applying the JAT.

Original: Although the performance gain decreases a little, we acquire an even higher performance than JAT in the fine-tuning manner (66.5 > 65.4).

Original: This shows the potential of building models with the JAT architecture.

Original: [Table]

Original: G Limitations / Broder Impact

Original: Our method can enable the Transformer model to capture high-order connections by applying the JAT techniques.

Original: It can potentially benefit many communities in society.

Original: For example, the large-scale pre-train model (GPT-3,

Original: T5, etc.) can use our method to enhance the model performance without more examples.

Original: It can save energy consumption during collecting examples and training models for complex tasks.

Original: As another example, the popular CLIP model can use our method to enhance its performance on various tasks, especially in a dangerous environment.

Original: Our method requires GPUs and high-speed networks to run.

Original: For underdeveloped regions such as rural areas, our method may not be feasible to use.

Original: Our method develops on the input data, which may have domain shift and population bias.

Original: Such bias may render the JAT to produce suboptimal results.

Original: If the method fails, the Transformer may generate incorrect outputs.

Original: In mission-critical applications, the usage of this method should

Original: [Table]

Original: [Table]