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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

Original: Abstract
Revised: Abstract

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

Original: However,the strong assumption that elements are directly attentive to each other limits theperformance 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 theperformance of tasks with high-order dependencies such as natural language understanding and Image captioning.

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

Original: Inspired by the pieces movingof English Draughts, we introduce the spectral convolutional technique to calculate
Revised: Inspired by the pieces movingof 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 ineach self-attention head and is interchangeable with the canonical self-attention.
Revised: This technique allows JAT’s propagation ineach self-attention head and is interchangeable with the canonical self-attention.

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

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

Original: With extensive experiments, we empirically show that ourmethods significantly increase the performance on ten different tasks
Revised: With extensive experiments, we empirically show that ourmethods 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 [4] and switch Transformer [11], with their massivelayers and over-parameterized architectures, have not further achieved significant performanceimprovements since the BERT model [8] beat the Human Performance in GLUE [31] and SQuAD [21].
Revised: However, the giant model GPT-3 [4] and switch Transformer [11], with their massivelayers and over-parameterized architectures, have not further achieved significant performanceimprovements since the BERT model [8] beat the Human Performance in GLUE [31] and SQuAD [21].

Original: Recently two visualization works [30, 6] reveal that the self-attention mechanism mainly learnssuperficial connections among inputs.
Revised: Recently two visualization works [30, 6] reveal that the self-attention mechanism mainly learnssuperficial connections among inputs.

Original: The principle of direct accessing among inputs in the dotproduct self-attention lies as the core ofthe above models, which solely relies on their mutualsimilarity 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 inthe vector space.

Original: However, as “every coin has two sides”, while the dot-product is simpleand effective in capturing long-range dependencies, it has limited capability in learning high-orderdependencies.
Revised: However, as “every coin has two sides”, while the dot-product is simple and effectivein 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 gamespread to Great Britain in about 1100.
Revised: To solve this problem, we are inspired by the pieces moving of English Draughts, a checker gamespread to Great Britain in about 1100.

Original: It allows two moving rules: one is the simple move from thecurrent position to an adjacent position; and the other is the capture move jumping over the “steppingstones” (opponent’s pieces) from front to back.
Revised: It allows two moving rules: one is the simple move from thecurrent position to an adjacent position; and the other is the capture move jumping over the “steppingstones” (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 movesin 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 movesin English Draughts, to enhance the feature map of dot-product self-attention.

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

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

Original: and “Bookcosts less than computer”.
Revised: and “Bookcosts 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 thisexample in detail in Section 3.1 and clarify the definition of Jump Self-attention.
Revised: We will discuss thisexample in detail in Section 3.1 and clarify the definition of Jump Self-attention

Original: [Figure]

Original: In this work, we find that the high-order self-attention exists in the dot-product computation but it istoo 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 cancapture such high-order dependencies, which is shown in Fig.(2) with further discussions.
Revised: Only a few heads in Transformers can [Figure] capture such high-order dependencies, which is shown in Fig.(2) with further discussions.

Original: Insteadof finding such "miracle" heads, we reformulate them to enhance the high-order dependencieswith the feature propagation technique in graph learning.
Revised: Insteadof finding such "miracle" heads, we reformulate them to enhance the high-order dependencieswith the feature propagation technique in graph learning.

Original: Meanwhile, we build an efficient variantin computing-sensitive situations based on the sparsity measurement [36], and build higher-ordervariants in comprehensive situations using MixHop [1].
Revised: Meanwhile, we build an efficient variantin computing-sensitive situations based on the sparsity measurement [36], and build higher-ordervariants 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 Transformermodels, 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 Transformermodels, where we formulate the high-order statistics calculation into a graph propagation problem.

Original: Itcan be used interchangeably with the canonical self-attention [29].
Revised: Itcan be used interchangeably with the canonical self-attention [29].

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

Original: Our main contributionscan be summarized as:
Revised: Our main contributionscan be summarized as:

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

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

Original: We propose the higher-order variants of JAT and masking strategy to allow its interchangeableusage 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 interchangeableusage with the canonical self-attention on the pre-trained models.

Original:  We conduct experiments on ten different tasks to examine the effectiveness of the proposedmethods.
Revised: • We conduct experiments on ten different tasks to examine the effectiveness of the proposedmethods.

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 [29] model for the encoder-decoder architecture in
Revised: The self-attention was introduced by Transformer [29] model for the encoder-decoder architecture in

Original: NLP to build long-range dependencies from a variable-length source sentence.
Revised: NLP to build long-range dependencies from a variable-length source sentence.

Original: The alignment abilityof self-attention to learn to acquire direct access between important items has made it an efficientcomponent in many applications [8, 19, 34].
Revised: The alignment abilityof self-attention to learn to acquire direct access between important items has made it an efficientcomponent in many applications [8, 19, 34].

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

Original: [Equation]
Revised: [Equation]

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

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

Original: In practice, multiple attention heads are used tolearn robust patterns [6, 30] among different transformations.
Revised: In practice, multiple attention heads are used tolearn robust patterns [6, 30] among different transformations.

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

Original: We concatenate the feature map of allgroups as the attention layer output.[Figure] [Figure]
Revised: We concatenate the feature map of allgroups 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.

Revised: [Figure]

Revised: [Figure]

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

Original: The BERT [8] model and its pre-trained successors utilize the self-attention, like GPT-3 [4] andswitch Transformer [11], and they achieve distinct performance improvements on various NLPtasks.
Revised: The BERT [8] model and its pre-trained successors utilize the self-attention, like GPT-3 [4] andswitch Transformer [11], and they achieve distinct performance improvements on various NLPtasks.

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

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

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

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

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

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

Original: This paperinvestigates whether it is worth enhancing lexical patterns in self-attention.
Revised: This paperinvestigates 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 patternsof self-attention heads.
Revised: To motivate our approach, we perform an individual case study on investigating the lexical patternsof self-attention heads.

Original: Based on the pre-trained model on BERT [8], we consider logical reasoningquestions on two continuous sentences “Book cost more than pencil” and “Book cost less thancomputer” and it seeks to find the price between “pencil” and “computer”.
Revised: Based on the pre-trained model on BERT [8], we consider logical reasoningquestions on two continuous sentences “Book cost more than pencil” and “Book cost less thancomputer”. It seeks to find the price between entities that have not been compared, like “pencil”and “computer”.

Original: Note that the articles areomitted 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-attentionscores at the seventh layer and fifth head.
Revised: In Fig.(2a), we present thevisualization of three words’ self-attention scores at the seventh layer and fifth head.

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

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

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

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

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

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

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

Original: We onlydraw the local connections among the three words and it covers 24.5% of the total as the significantlocal connections.
Revised: Keys [Figure] draw the local connections among the three words and it covers 24.5% of the total as the significantlocal connections.

Original: The stronger “Computer-Pencil” attention connection, as it in Fig(2b), is the soleself-attention computation that happens at the seventh 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 soleself-attention computation that happens at the seventh layer, which distinguishes it from the trivial“Book-pencil” and “Book-Computer” attention connections.

Original: As we discussed above, there are manyconnections from three words to the delimiter ‘CLS’ and positional tokens ‘SEP’ but they did notform the “Computer-Pencil” connection.
Revised: As we discussed above, there are manyconnections from three words to the delimiter ‘CLS’ and positional tokens ‘SEP’ but they did notform 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 capturethe “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 capturethe “Computer-Pencil” connection.

Original: The fifth head in the seventh layer is the miracle head that canestablish the strong linguistic connection between “computer” and “Pencil” and it forms lexicalself-attention patterns across different sentences.
Revised: The fifth head in the seventh layer is the miracle head that canestablish the strong linguistic connection between “computer” and “Pencil” and it forms lexicalself-attention patterns across different sentences.

Original: We call this an exact example of the High-order
Revised: We call this an exact example of the High-order

Original: Statistics , and it is formulated on input pairs which satisfy the following definition:
Revised: 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 singlehead 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 singlehead 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 connectionscovering 80.2% local connections among the three words.
Revised: In Fig.(2c), we relax the restriction on significant connections and select the major connectionscovering 80.2% local connections among the three words.

Original: As a natural result of decreasing thesignificant δ 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 thesignificance δ 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-attentionsexist 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-attentionsexist 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 withthe latest improvement on it.
Revised: The Jump Self-attention (JAT) is based on the vanilla Transformer architecture to be consistent withthe latest improvement on it.

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

Original: Then JAT utilizes the GCN operator to allow “jump” connections totraverse through different queries and key pairs, and the weak connections under the similar conditionare enhanced to help enforces the desired triadic closure property.
Revised: Then JAT utilizes the GCN operator to allow “jump” connections totraverse through different queries and key pairs, and the weak connections under the similar conditionare 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 thesoftmax( · ) 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 thesoftmax( · ) 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 attentioncomputation, we can take a column view of S in Fig.(3a).
Revised: To understand the underlying meaning of high-order statistics in attentioncomputation, we can take a column view of S in Fig.(3a).

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

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

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

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 vectorof the first key, as we can see the crosshatch under the column view (the red shade), which indicatesthat 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 vectorof the first key, as we can see the crosshatch under the column view (the red shade), which indicatesthat inputs x 3 and x 4 attend to the input x 1 from the K 1 ’s observation.

Original: The connection between x 3and x 4 are weak (the blank square).
Revised: The connection between x 3and 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 inputx 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 inputx 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 attainslarger attention scores while others maintains, and the attention score between x 3 and x 4 grows largerenough than the significant threshold p .
Revised: If we could modify the linear transformation W Q and W K such that the red marked square attainslarger attention scores while others maintains, and the attention score between x 3 and x 4 grows largerenough than the significance threshold p .

Original: We call this attention adjustment a successful “jump”, and itcould happen for the individual key.
Revised: We call this attention adjustment a successful “jump”, andit 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

Original: Queries

Original: Keys

Original: Queries

Original: Keys

Original: Queries

Original: Keys

Original: Queries

Original: [Figure]

Original: [Figure]

Original: Q 1 : Book

Original: Q 4 : Computer Q 3 : Pencil

Original: K 1 : Book K 3 : Pencil K

Original: [Figure]

Original: Computer

Original: Q 1 : Book

Original: Q 4 : Computer Q 3 : Pencil

Original: K 1 : Book K 3 : Pencil K 4 : Computer

Original: Q 1 : Book

Original: Q 4 : Computer Q 3 : Pencil

Original: K 1 :

Original: [Figure]

Original: K

Original: [Figure]

Original: K 4 : Computer +

Original: [Figure]

Original: Q 1 :

Original: [Figure]

Original: Q

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

Original: However, this direct operation will introduce an extracomputation of attention combination O
Revised: However, this direct operation will introduce an extracomputation of attention combination 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 aback-propagation friendly implementation.
Revised: And it is hard to acquire aback-propagation friendly implementation.

Original: To formulate the transformation modification, we relax the objective of finding the miracle heads toenhance 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 toenhance 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 in the set E is defined by the adjacencymatrix A .
Revised: V with the row vector of S as node features and each edge in the set E is defined by the adjacencymatrix A .

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 keyas the current column, then the sub-adjacency matrix under the K j is defined:
Revised: We select the j -th keyas the current column, then the sub-adjacency matrix under the K j is defined:

Original: [Equation]
Revised: [Equation]

Original: where U j = S j ⊤ ⊗ S j ⊤ 
Revised: where U j = S j ⊤ ⊗ S j ⊤ d .

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

Original: U j is a symmetric matrix contains 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 itemthat returns 1 if the item is greater than the ρ , otherwise return 0 .
Revised: [ · ] ρ is an indicator function for each matrixitem that returns 1 if the item is greater than the ρ , otherwise return 0 .

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

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

Original: We go through all keys, the j -th sub-adjacencymatrix 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 similaritybetween different queries w.r.t K j .

Original: Inspired by the layer aggregation [28, 18], the summation of weighted adjacency matrix yieldlarger gap between isolated values and bulk eigenvalues, which has been widely used in helpingthe community detection task [24, 7].
Revised: Inspired by the layer aggregation [28, 18], the summation of weighted adjacency matrix yieldlarger gap between isolated values and bulk eigenvalues, which has been widely used in helpingthe community detection task [24, 7].

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

Original: Then, we can add up all thesub-adjacency matrix into the weighted one:
Revised: Then, we can add up all thesub-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: [25] and is further simplified by Duvenaud et al.
Revised: [25] and is further simplified by Duvenaud et al.

Original: [10] and
Revised: [10] and

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

Original: We employ it as a specific architecture for encourage the attended itemsbecomes similar in both queries and keys.
Revised: We employ it as a specific architecture for encouraging the attended items tobecome 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 whenencountering each other in self-attention computation.
Revised: GCN will make the selected queries/keys propagate its features, which causes a higher score whenencountering each other in self-attention computation.

Original: In other words, GCN encourages high-orderself-attention’s propagation through individual keys’ similarity measurement.
Revised: In other words, GCN encourages high-orderself-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 ′ , thescore 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 ′ , thescore S ′ becomes ˆ A ( XW ′ Q )( XW ′ K )

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

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 1 heads and H 2 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 1 heads and H 2 heads.

Original: The H 1 heads perform the canonical selfattention operation and the H 2 heads perform high-order self-attention, then we concatenate theirresults as the complete layer output.
Revised: The H 1 heads perform the canonical selfattention operation and the H 2 heads perform high-order self-attention, then we concatenate theirresults as the complete layer output.

Original: The use of high-order self-attention enables acquiring high-orderdependencies over the superficial ones.
Revised: The use of high-order self-attention enables acquiring high-orderdependencies 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 [29] plays an important role in avoiding the auto-regressive during theinference phase.
Revised: The mask of self-attention [29] plays an important role in avoiding the auto-regressive during theinference 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 wereduce 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 wereduce the edges from afterward queries to previous queries.

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

Original: The canonical self-attention’s maskingoperation on the feature map is preserved after applying GCN in Eq.(5) such that we can eschew theinformation leakage in softmax( · ) function.
Revised: The canonical self-attention’s maskingoperation on the feature map is preserved after applying GCN in Eq.(5) such that we can eschew theinformation 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 totraverse 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 totraverse all keys in the adjacency matrix aggregation of Eq.(3).

Original: Zhou et al.
Revised: Zhou et al.

Original: [36] design a sparsitymeasurement to find the most significant queries in self-attention, which motivates us to employ it inselecting the primary keys:
Revised: [36] design a sparsitymeasurement to find the most significant queries in self-attention, which motivates us to employ it inselecting 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 reducesto ¯ S ∈ R L × u and we acquire the efficient variant JAT † in the O
Revised: Thus, the original self-attention S reducesto ¯ S ∈ R L × u and we acquire the efficient variant JAT † in the O

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

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

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

Original: Without loss of generality, we can extend the similar in Definition 3.1 to a higher order case, wherethe similarity depends on more than one input.
Revised: Without loss of generality, we can extend the similar in Definition 3.1 to a higher order case, wherethe similarity depends on more than one input.

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

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

Original: [Equation]
Revised: [Equation]

Original: [Table]

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 tosubstitute the Φ j ( S ) in Eq.(7) for higher-order self-attention variants A2 , A 3 and etc.
Revised: It’s trivial tosubstitute the Φ j ( S ) in Eq.(7) for higher-order self-attention variants A

Original: We can splitthe heads into more groups when concatenating multiple self-attentions with different orders.
Revised: [Table]

Original:  Experiments
Revised:  Experiments

Original:  In this section, we empirically demonstrate JAT’s effectiveness on the pre-trained BERT base modelwith 12 layers, and perform BERT-JAT fine-tuning and evaluation on ten NLP tasks.
Revised: [Table] In this section, we empirically demonstrate JAT’s effectiveness on the pre-trained BERT base modelwith 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[31], a collection ofdiverse 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[31], a collection ofdiverse natural language understanding tasks.

Original: Settings: We use a batch size of 32 and fine-tune for 5epochs over the data for nine GLUE tasks.
Revised: Settings: We use a batch size of 32 and fine-tune for 5epochs 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: Theother settings follow the recommendation of the original paper.
Revised: Theother settings follow the recommendation of the original paper.

Original: Since the proposed JAT can be usedinterchangeably with canonical self-attention, we perform a grid search on the layer replacement.
Revised: Since the proposed JAT can be usedinterchangeably 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 isthe 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 isthe 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 the9 tasks.
Revised: Metric: We use three different evaluation metrics on the9 tasks.

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

Original: [Table] √( [Table] [Equation]
Revised: TP × TN − FP × FN √

Original: [Table] [Table] + [Equation]
Revised: ( TP + FP )

Original:  )( TN [Equation] FP )( TN + FN ) ) is usedfor CoLA, Pearson ( PCC =[Equation]
Revised: ( TP + FN )( TN + FP )( TN + FN ) ) is usedfor CoLA, Pearson ( PCC =

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table] [Equation] [Table] [Equation]
Revised: √

Original: [Table]

Original: [Equation]

Original: N
Revised: N

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table] [Equation] [Table] [Equation]
Revised: √ N

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: [Table]
Revised: [Table]

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

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

Original: (b) SQuAD. We also evaluate on the SQuAD v1.1/v2.0 [21, 22].
Revised: We also evaluate on the SQuAD v1.1/v2.0 [21, 22].

Original: The settings is in Appendix C.
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 [13] the baselineof GLUE, BERT base [8], and RoBERTa [17], the pre-trained model where we build JAT varints.
Revised: We compare ELMo [13] the baselineof GLUE, BERT base [8], and RoBERTa [17], the pre-trained model where we build JAT varints.

Original: Ourmodel BERT/RoBERTa-JAT outperform base models on 8 tasks without introducing more parameters,which shows that JAT can capture higher-order dependencies than canonical self-attention.
Revised: Ourmodel BERT/RoBERTa-JAT outperform base models on 8 tasks without introducing more parameters,which shows that JAT can capture higher-order dependencies than canonical self-attention.

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

Original: When it applies to RoBERTa, the average gain is larger and the most significant rising 3.2%happens on the RTE 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 competitor A 3 [35] generally.
Revised: We have also show that JAT beats competitor A 3 [35] generally.

Original: The A 3 introduces the tripletattention measuring dissimilarity to diversify the self-attention, while our proposed JAT tried toenhance similarity by building connections between unrelated pairs from a different perspective.
Revised: The A 3 introduces the tripletattention measuring dissimilarity to diversify the self-attention, while our proposed JAT tried toenhance 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: The result of the SQuAD experiment is summarized in Table 2.
Revised: The result of the SQuAD experiment is summarized in Table 2.

Original: Our model BERT-JAT achieves betterperformance than all baselines in both EM and F1 scores, which is aligned with the GLUE’s results.
Revised: Our model BERT-JAT achieves betterperformance than all baselines in both EM and F1 scores, which is aligned with the GLUE’s results.

Original: [Table]
Revised: 4.3.1 Effect of JAT’s Layer Deployment In this experiment, we explore the JAT layer’s effect on acquiring the high-order connections. Weset the layers at a different position in RoBERTa base model to use JAT heads, where the RoBERTa-

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

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

Original: In this experiment, we explore the JAT layer’s effect on acquiring the high-order connections. Weset the layers at a different position in BERT base model to use JATheads, where the BERT-JAT l [ i,j ]represents that the layers from i to j in BERT base model have four JAT heads.
Revised: JAT l [ i,j ] represents that the layers from i to j in RoBERTa base model have four JAT heads.

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

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

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

Original: Applying JAT at lower layers can effectively enhance the diversity of lower layers and improve thehigh-order connections to force the desired triadic closure.
Revised: Applying JAT at lower layers can effectively enhance the diversity of lower layers andimprove 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 taskaccuracy, we trained many BERT-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 understandingtask accuracy, we trained many RoBERTa-JAT models with different JAT head numbers in layer 1-6.

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

Original: With the number of JAT head increasing, we canfind 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 BERT model.
Revised: The scores increase becauseadding JAT heads enhance the high-order global dependencies in the vanilla RoBERTa model.

Original: They decrease when thenumber of heads is greater than 6, which probably is caused by too strong high-order inductive biassuppressing the canonical self-attention.
Revised: Theydecrease 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: But the CoLA dataset seems fitting the

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 similarsettings as the previous ablation study, where the BERT-JAT oi stands for the i -th order self-attentiondefined in Eq.(9).
Revised: Following the similarsettings 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 5.
Revised: We show the selected GLUE results in Table 5.

Original: And it reveals that for mostnatural language understanding tasks, the second-order self-attention is sufficient to discover morehigher-order connections and improve the expression ability of the model, which leads to betterresults.
Revised: And it reveals that formost natural language understanding tasks, the second-order self-attention is sufficient to discovermore higher-order connections and improve the expression ability of the model, which leads to betterresults.

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.4 Visualization
Revised: 4.4 Visualization

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

Original: To avoidrepetition, we construct a new sentence.
Revised: To avoidrepetition, we construct a new sentence.

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

Original: The Fig.(4a) 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.(4a) 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.(4b) shows heads of the samelayer in BERT-JAT, the “Kate” attends to the first sentence’s “hit” through a “jump” connection withthe blue head.
Revised: The Fig.(4b) shows heads of the samelayer in BERT-JAT, the “Kate” attends to the first sentence’s “hit” through a “jump” connection withthe 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’), whichindicates 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’), whichindicates the two actions may be sequentially related.

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

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

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

Original: [Figure]
Revised: [Figure]

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

Original: 4.5 Case Study
Revised: 4.5 Case Study

Original: We conduct an experiment on Named-Entity Recognition (NER) task using the standard CoNLL-2003dataset, which concentrates on four types of named entities: persons, locations, organizations andnames 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 CoNLL-2003dataset, which concentrates on four types of named entities: persons, locations, organizations andnames of miscellaneous entities that do not belong to the previous three groups.

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

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

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

Original: From Fig.(6a) we find there are notmany ‘CROSS’ connections in the attention of BERT base , and most of the connections are betweentokens belong to person entity.
Revised: From Fig.(6a) we find there are notmany ‘CROSS’ connections in the attention of BERT base , and most of the connections are betweentokens belong to person entity.

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

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 hasno learnable parameters.
Revised: With the same setting, we evaluate the iterative efficiency of JAT on eight downstream task, which hasno learnable parameters.

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

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

Original: It could be further improved byutilizing the Keys’ sparsity [36], then JAT † achieves no more than 15% extra computation.
Revised: It could be further improved byutilizing the Keys’ sparsity [36], then JAT † achieves no more than 15% extra computation.

Original:  Conclusions
Revised:  Conclusions

Original: In this paper, we discovered the existence of Jump Self-attention (JAT) in the Transformers’ dotproduct computation and formalized the capturing high-order statistics into the graph learningproblem with multi-hop variants.
Revised: In this paper, we discovered the existence of Jump Self-attention (JAT) in the Transformers’ dotproduct computation and formalized the capturing high-order statistics into the graph learningproblem with multi-hop variants.

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

Original: Extensive experiments demonstrate JAT’s superior performance on two tasks.
Revised: Extensive experiments demonstrate JAT’s superior performance on two tasks.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu-
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu-

Original: tions and scope?
Revised: tions and scope?

Original: [Yes] Please refer to Sec.1.(b)
Revised: [Yes] Please refer to Sec.1.

Original:  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 Broad ImpactSection.(c)
Revised: [Yes] We include that in the Broad ImpactSection.(c)

Original: Did you discuss any potential negative societal impacts of your work?
Revised: Did you discuss any potential negative societal impacts of your work?

Original: [Yes] We include that in
Revised: [Yes] We include that in

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

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes]
Revised: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

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

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experimental
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experimental

Original: results (either in the supplemental material or as a URL)?
Revised: results (either in the supplemental material or as a URL)?

Original: [Yes] We provide all of them in the
Revised: [Yes] We provide all of them in the

Original: supplemental material.
Revised: supplemental material.

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
Revised: [Yes] Please

Original: [Table]
Revised: [Table]

Original: Sec.4.
Revised: to Sec.4.1.

Original: [Table]

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

Original: [Table]
Revised: [Table]

Original: multiple times)?
Revised: multiple times)?

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

Original: (d) Did you include the total amount of compute and the type of resources used (e [Table] g., type of GPUs,
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs,

Original: [Table]
Revised: [Table]

Original: internal cluster, or cloud provider)?
Revised: 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
Revised: If you are using existing assets (e

Original: [Table]
Revised: [Table]

Original: g., code, data, models) or curating/releasing new assets...
Revised: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re us-
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re us-

Original: ing/curating?
Revised: ing/curating?

Original: [Yes]
Revised: [Yes]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: informa-
Revised: informa-

Original: tion or offensive content?
Revised: tion 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...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: [Table]
Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]