B I -D IRECTIONAL B LOCK S ELF -A TTENTION FOR F AST AND M EMORY -E FFICIENT S EQUENCE M ODELING
A BSTRACT
Recurrent neural networks (RNN), convolutional neural networks (CNN) and selfattention networks (SAN) are commonly used to produce context-aware representations.
RNN can capture long-range dependency but is hard to parallelize and not time-efﬁcient.
CNN focuses on local dependency but does not perform well on some tasks.
SAN can model both such dependencies via highly parallelizable computation, but memory requirement grows rapidly in line with sequence length.
In this paper, we propose a model, called “bi-directional block self-attention network (Bi-BloSAN)”, for RNN/CNN-free sequence encoding.
It requires as little memory as RNN but with all the merits of SAN.
Bi-BloSAN splits the entire sequence into blocks, and applies an intra-block SAN to each block for modeling local context, then applies an inter-block SAN to the outputs for all blocks to capture long-range dependency.
Thus, each SAN only needs to process a short sequence, and only a small amount of memory is required.
Additionally, we use feature-level attention to handle the variation of contexts around the same word, and use forward/backward masks to encode temporal order information.
On nine benchmark datasets for different NLP tasks, Bi-BloSAN achieves or improves upon state-of-the-art accuracy, and shows better efﬁciency-memory trade-off than existing RNN/CNN/SAN.
I NTRODUCTION
Context dependency provides critical information for most natural language processing (NLP) tasks.
In deep neural networks (DNN), context dependency is usually modeled by a context fusion module, whose goal is to learn a context-aware representation for each token from the input sequence.
Recurrent neural networks (RNN), convolutional neural networks (CNN) and self-attention networks (SAN) are commonly used as context fusion modules.
However, each has its own merits and defects, so which network to use is an open problem and mainly depends on the speciﬁc task.
RNN is broadly used given its capability in capturing long-range dependency through recurrent state computation.
It has been applied to various NLP tasks, e.g., question answering (Wang et al., 2017), neural machine translation (Bahdanau et al., 2015), sentiment analysis (Li et al., 2017), natural language inference (Liu et al., 2016), etc.
However, training the basic RNN encounters the gradient dispersion problem, and is difﬁcult to parallelize.
Long short-term memory (LSTM) (Hochreiter &
Schmidhuber, 1997) effectively avoids the vanishing gradient.
Gated recurrent unit (GRU) (Chung et al., 2014) and simple recurrent unit (SRU) (Lei & Zhang, 2017) improve the efﬁciency by reducing parameters and removing partial temporal-dependency, respectively.
However, they still suffer from expensive time cost, especially when applied to long sequences.
CNN becomes popular recently on some NLP tasks because of its the highly parallelizable convolution computation.
Unlike RNN, CNN can simultaneously apply convolutions deﬁned by different kernels to multiple chunks of a sequence (Kim, 2014).
It is mainly used for sentence-encoding tasks (Lei et al., 2015; Kalchbrenner et al., 2014).
Recently, hierarchical CNNs, e.g. ByteNet (Kalchbrenner et al., 2016), and ConvS2S (Gehring et al., 2017), are proposed to capture relatively long-range dependencies by using stacking CNNs to increase the number of input elements represented in a state.
Nonetheless, as mentioned by Vaswani et al.
(2017), the number of CNNs required to relate signals from two arbitrary input grows in the distance between positions, linearly for ConvS2S and logarithmically for ByteNet.
This makes it difﬁcult to learn dependencies between distant positions.
Recently, self-attention networks (SAN) have been successfully applied to several NLP tasks.
It produces context-aware representation by applying attention to each pair of tokens from the input sequence.
Compared to RNN/CNN, SAN is ﬂexible in modeling both long-range and local dependencies.
The major computation in SAN is the highly parallelizable matrix multiplication without any temporal iteration, which can be easily accelerated by existing tools.
Unlike most works that attach SAN to RNN/CNN as an additional module, two recent works show that SAN independent of any RNN/CNN module can achieve state-of-the-art performance on several NLP tasks.
The ﬁrst, multi-head attention (Vaswani et al., 2017), is a major component of a seq2seq model “Transformer” that outperforms previous methods in neural machine translation.
It projects the input sequence into multiple subspaces, applies a SAN to the representation in each subspace, and concatenates the outputs.
The second, directional self-attention network (DiSAN) (Shen et al., 2017), computes alignment scores at feature level, rather than at token level, and applies forward/backward masks to the alignment score matrix to encode temporal order information.
DiSAN achieves the best or state-ofthe-art test accuracy on several NLP tasks by using less computational time and fewer parameters.
More related works can be found in Appendix D.
However, one drawback of SAN is its large memory requirement to store the alignment scores of all the token pairs; the number grows quadratically with the sequence length.
By contrast, RNN/CNN demand far less memory.
The goal of this paper is to develop a novel SAN for RNN/CNN-free sequence encoding, which requires as little memory as RNN but inherits all the advantages of SAN, i.e., highly parallelizable computation, the
We propose an attention mechanism, called “bidirectional block self-attention (Bi-BloSA)”, for fast and memory-efﬁcient context fusion.
The basic idea is to split a sequence into several length-equal blocks (with padding if necessary), and apply an intra-block
SAN to each block independently.
The outputs for all the blocks are then processed by an inter-block
SAN.
The intra-block SAN captures the local dependency within each block, while the inter-block SAN captures the long-range/global dependency.
Hence, every SAN only needs to process a short sequence
Compared to a single SAN applied to the whole sequence, such two-layer stacked SAN saves a signiﬁcant amount of memory.
A feature fusion gate combines the outputs of intra-block and inter-block SAN with the original input, to produce the ﬁnal contextaware representations of all the tokens.
Similar to directional self-attention (DiSA) (Shen et al., 2017), BiBloSA uses forward/backward masks to encode the temporal order information, and feature-level attention to handle the variation of contexts around the same word.
Further, a RNN/CNN-free sequence encoding model we build based on Bi-BloSA, called “bi-directional block self-attention network (Bi-BloSAN)”, uses an attention to compress the output of Bi-BloSA into a vector representation.
In experiments 1 , we implement Bi-BloSAN and popular sequence encoding models on several NLP tasks, e.g., language inference, sentiment analysis, semantic relatedness, reading comprehension, question-type classiﬁcation, etc.
The baseline models include Bi-LSTM, Bi-GRU, Bi-SRU, CNNs, multi-head attention and DiSAN.
A thorough comparison on nine benchmark datasets demonstrates the advantages of Bi-BloSAN in terms of training speed, inference accuracy and memory consumption.
Figure 1 shows that Bi-BloSAN obtains the best accuracy by costing similar training time to DiSAN, and as little memory as Bi-LSTM, Bi-GRU and multi-head attention.
This shows that
Bi-BloSAN achieves a better efﬁciency-memory trade-off than existing RNN/CNN/SAN models.
Our notations follow these conventions: 1) lowercase denotes a vector; 2) bold lowercase denotes a sequence of vectors (stored as a matrix); and 3) uppercase denotes a matrix or a tensor.
B ACKGROUND 2. W ORD E MBEDDING
Word embedding is the basic processing unit in most DNN for sequence modeling.
It transfers each discrete token into a representation vector of real values.
Given a sequence of tokens (e.g., words or characters)
w =
[ w 1 , w 2 , . .
w n ] ∈ R N × n , where w i is a one-hot vector, N is the vocabulary size and n is the sequence length.
A pre-trained token embedding (e.g. word2vec (Mikolov et al., 2013b)) is applied to w , which outputs a sequence of low dimensional vectors x = [ x 1 , x 2 , . . .
, x n ] ∈ R d e × n .
This process can be formally written as x = W ( e ) w , where
W ( e ) ∈ R d e × N is the embedding weight matrix that can be ﬁne-tuned during the training phase.
2. V ANILLA A TTENTION AND M ULTI - DIMENSIONAL A TTENTION
Vanilla Attention :
Given an input sequence x = [ x 1 , x 2 , . . .
, x n ] composed of token embeddings and a vector representation of a query q ∈ R d q , vanilla attention (Bahdanau et al., 2015) computes the alignment score between q and each token x i (reﬂecting the attention of q to x i ) using a compatibility function f
( x i , q ) .
A softmax function then transforms the alignment scores a ∈ R n to a probability distribution p ( z | x , q ) , where z is an indicator of which token is important to q .
A large p ( z = i | x , q ) means that x i contributes important information to q .
This process can be written as
[Equation]
The output s is the expectation of sampling a token according to its importance, i.e.,
[Equation]
Multiplicative attention (or dot-product attention) (Vaswani et al., 2017; Sukhbaatar et al.,
[Equation]
Rush et al., 2015) and additive attention (or multi-layer perceptron attention) (Bahdanau et al., 2015;
Shang et al., 2015) are two commonly used attention mechanisms.
They differ in the choice of compatibility function f
( x i , q ) .
Multiplicative attention uses the cosine similarity for f ( x i , q ) , i.e.,
[Equation]
where W (1) ∈ R d h × d e , W (2) ∈ R d h × d q are the learnable parameters.
Additive attention is deﬁned as
[Equation]
where w ∈ R d h , b (1) and b are the biases, and σ ( · ) is an activation function.
Additive attention usually achieves better empirical performance than multiplicative attention, but is expensive in time cost and memory consumption.
Multi-dimensional Attention : Unlike vanilla attention, in multi-dimensional (multi-dim) attention (Shen et al., 2017), the alignment score is computed for each feature, i.e., the score of a token pair is a vector rather than a scalar, so the score might be large for some features but small for others.
Therefore, it is more expressive than vanilla attention, especially for the words whose meaning varies in different contexts.
Multi-dim attention has d e indicators z 1 , . . .
, z d e for d e features.
Each indicator has a probability distribution that is generated by applying softmax to the n alignment scores of the corresponding feature.
Hence, for each feature k in each token i , we have P ki (cid:44) p ( z k
= i | x , q ) where P ∈ R d e × n .
A large P ki means that the feature k in token i is important to q .
The output of multi-dim attention is written as
[Equation]
For simplicity, we ignore the subscript k where no confusion is caused.
Then, Eq.(6) can be rewritten as an element-wise product, i.e., s = (cid:80) ni =1 P · i (cid:12) x i .
Here, P · i is computed by the additive attention in Eq.(5) where w T is replaced with a weight matrix W ∈ R d h × d e , which leads to a score vector for each token pair.
2. T WO T YPES OF S ELF - ATTENTION
token2token self-attention (Hu et al., 2017; Vaswani et al., 2017; Shen et al., 2017) produces context-aware representations by exploring the dependency between two tokens x i and x j from the same sequence x .
In particular, q in Eq.(5) is replaced with x j , i.e.,
[Equation]
Similar to the P in multi-dim attention, each input token x j is associated with a probability matrix
P j such that P jki (cid:44) p ( z k = i | x , x j ) .
The output representation for x j is s j = (cid:80) n i =1 P j · i (cid:12) x i and the ﬁnal output of token2token self-attention is s =
[ s 1 , s 2 , . .
, s n ] .
source2token self-attention (Lin et al., 2017; Shen et al., 2017; Liu et al., 2016) explores the importance of each token to the entire sentence given a speciﬁc task.
In particular, q is removed from
Eq.(5), and the following equation is used as the compatibility function.
[Equation]
The probability matrix P is deﬁned as P ki (cid:44) p ( z k = i
| x ) .
The ﬁnal output of source2token selfattention has the same form as multi-dim attention, i.e., s = (cid:80) ni
=1 P · i (cid:12) x i
2. M ASKED S ELF - ATTENTION
Temporal order information is difﬁcult to encode in token2token self-attention introduced above because the alignment score between two tokens is symmetric.
Masked self-attention (Shen et al., 2017) applies a mask M ∈ R n × n to the alignment score matrix (or tensor due to feature-level score) computed by Eq.(7), so it allows one-way attention from one token to another.
Speciﬁcally, the bias b in Eq.(7) is replaced with a constant vector M ij 1 , where the 1 is an all-one vector.
In addition, W is ﬁxed to a scalar c and tanh( · /c ) is used as the activation function σ ( · ) , i.e.,
[Equation]
where W (1) ∈ R d e × d e , W (2) ∈ R
d e × d q .
The procedures to calculate the attention output from f ( x i , x j ) are identical to those in token2token self-attention.
We use s = g m ( x , M ) to denote the complete process of masked self-attention with s =
[ s 1 , s 2 , . . .
, s n ] as the output sequence.
An illustration of masked self-attention is given in Figure 2.
In order to model bi-directional order information, forward mask M fw and backward mask M bw are respectively substituted into Eq.(9), which results in forward and backward self-attentions.
These two masks are deﬁned as
[Equation]
The outputs of forward and backward self-attentions are denoted by s fw
= g m ( x , M fw ) and s bw =
g m ( x , M bw ) , respectively.
P ROPOSED M ODEL
In this section, we ﬁrst introduce the “masked block self-attention (mBloSA)”
(Section 3.1) as a fundamental self-attention module.
Then, we present the “bi-directional block self-attention network (Bi-BloSAN)”
(Section 3.2) for sequence encoding, which uses the “bi-directional block selfattention (Bi-BloSA)”
(mBloSA with forward and backward masks) as its context fusion module.
3. M ASKED B LOCK S ELF - ATTENTION
As shown in Figure 3, masked block self-attention (mBloSA) has three parts from its bottom to top, i.e., 1) intra-block self-attention, 2) inter-block self-attention, and 3) the context fusion.
Intra-block self-attention : We ﬁrstly split the input sequence of token/word embeddings into m blocks of equal length r , i.e., [ x l ] ml =
=
[ x 1 , x 2 , . . .
, x m ] where x 1 = [ x 1 , x 2 , . . .
, x r ] , x 2 = [ x r +1 , x r +2 , . . .
, x 2 r ] and x m =
[ x n − r
+1 , x n − r +2 , . . .
, x n ] .
Padding can be applied to the last block if necessary.
Intra-block self-attention applies the masked self-attentions g m ( · , M ) with shared parameters to all the blocks , i.e.,
[Equation]
Its goal is to capture the local context dependency inside each block.
Similar to x l , the output representations of the tokens in the l -th block are denoted by h l =
[ h r ( l − 1)+1 , h r ( l − 1)+2 , . .
, h r × l ] .
Note, the block length r is a hyper-parameter and m = n/r .
In Appendix A, we introduce an approach to selecting the optimal r , which results in the maximum memory utility rate in expectation.
Inter-block self-attention : To generate a vector representation v l of each block, a source2token self-attention g s 2 t ( · ) is applied to the output h l of the intra-block self-attention on each block, i.e.,
[Equation]
Note we apply the parameter-shared g s 2 t ( · ) to h l for different blocks.
This provides us with a sequence v =
[ v 1 , v 2 , . .
. , v m ] of local-context representations at block level.
Inter-block self- attention then applies a masked self-attention to v in order to capture the long-range/global dependency among the blocks, i.e.,
[Equation]
To combine the local and global context features at block level, a gate is used to merge the input and the output of the masked self-attention dynamically.
This is similar to the gates in LSTM.
The output sequence e = [ e 1 , . . .
, e m ] of the gate is computed by
[Equation]
Context fusion :
Given the long-range context representations e = [ e 1 , . . .
e m ] ∈ R d e × m at block level, we duplicate e l for r times to get e l =
[ e l , e l , . . .
, e l ] (each token in block l has the global context feature representation e l ).
Let E (cid:44) [ e l ] ml =1 ∈ R
d
e × n .
Now, we have the input sequence x of word embeddings, the local context features h produced by intra-block self-attention, and the long-range/global context features E produced by inter-block self-attention.
A feature fusion gate (Gong & Bowman, 2017) is employed to combine them, and generates the ﬁnal context-aware representations of all tokens, i.e.,
[Equation]
where σ ( · ) is an activation function, and u = [ u 1 , u 2 , . .
. , u n ] ∈ R d e × n is the mBloSA output, which consists of the context-aware representations of the n tokens.
3. B I - DIRECTIONAL B LOCK S ELF -A TTENTION N ETWORK FOR S EQUENCE ENCODING
We propose a sequence encoding model “Bi-directional block self-attention network (Bi-BloSAN)” with mBloSA as its major components.
Its architecture is shown in Figure 4.
In Bi-BloSAN, two fully connected layers (with untied parameters) are applied to the input sequence of token embeddings.
Their outputs are processed by two mBloSA modules respectively.
One uses the forward mask M fw and another uses the backward mask M bw .
Their outputs u fw and u bw are concatenated as u bi =
[ u fw ; u bw ] ∈ R 2 d e × n .
The idea of bi-directional attention follows the same spirit as Bi-LSTM and DiSAN.
It encodes temporal order information lacking in existing SAN models.
The context fusion module in
Bi-BloSAN, with the input x and the output u bi , is called “Bi-BloSA”.
In order to obtain a sequence encoding, a source2token self-attention transforms the sequence u bi of concatenated token representations into a vector representation s .
E XPERIMENTS
We conduct the experiments of Bi-BloSAN and several popular RNN/CNN/SAN-based sequence encoding models on nine benchmark datasets for multiple different NLP tasks.
Note that, in some baseline models, a source2token self-attention is on the top of the models to generate an encoding for the entire sequence.
All the models used for comparisons are listed as follows.
• Bi-LSTM :
600D Bi-directional LSTM (300D forward LSTM + 300D backward LSTM)
(Graves et al., 2013).
• Bi-GRU : 600D Bi-directional GRU (Chung et al., 2014).
• Bi-SRU :
600D Bi-directional SRU (Lei & Zhang, 2017) (with sped-up recurrence but no CUDA level optimization for fair comparison).
• Multi-CNN : 600D CNN sentence embedding model (Kim, 2014)
(200D for each of 3 , 4 , 5 -gram).
• Hrchy-CNN : 3-layer 300D CNN (Gehring et al., 2017) with kernel length 5, to which gated linear units (Dauphin et al., 2016) and residual connection
He et al., 2016) are applied.
• Multi-head :
600D Multi-head attention (Vaswani et al., 2017) ( 8 heads, each has 75 hidden units).
The positional encoding method used in Vaswani et al.
(2017) is applied to the input sequence to encode temporal order information.
DiSAN : 600D Directional self-attention network (Shen et al., 2017) (300D forward masked self-attention + 300D backward masked self-attention).
All experimental codes are implemented in Python with Tensorﬂow and run on a single Nvidia GTX 1080Ti graphic card.
Both time cost and memory load data are collected under Tensorﬂow1.3 with
CUDA8 and cuDNN6021.
In the rest of this section, we conduct the experiments on natural language inference in Section 4.1, reading comprehension in Section 4.2, semantic relatedness in Section 4.
and sentence classiﬁcations in Section 4.4.
Finally, we analyze the time cost and memory load of the different models vs. the sequence length in Section 4.5.
4. N ATURAL L ANGUAGE I NFERENCE
Natural language inference (NLI) aims to reason the semantic relationship between a pair of sentences, i.e., a premise sentence and a hypothesis sentence.
This relationship could be entailment , neutral or contradiction .
In the experiment, we compare Bi-BloSAN to other baselines on the Stanford Natural Language Inference (Bowman et al., 2015) (SNLI) 2 dataset, which contains standard training/dev/test split of 549,367/9,842/9,824 samples
[Table] [Table]
Following the method of applying sentence-encoding to NLI task given in Bowman et al.
[Table]
two parameter-tied sentence-encoding models are applied to the premise and the hypothesis sentences respectively, to generate the premise encoding s p and the hypothesis encoding s h .
A relation representation concatenating s p , s h , s p − s h and s p (cid:12) s h is passed into a 300D fully connected layer, whose output is given to a 3-unit output layer with softmax to calculate the probability distribution over the three classes.
Training Setup :
The optimization objective is the cross-entropy loss plus L2 regularization penalty.
We minimize the objective by Adadelta (Zeiler, 2012) optimizer which is empirically more stable than Adam (Kingma & Ba, 2015) on SNLI.
The batch size is set to 64 for all methods.
The training phase takes 50 epochs to converge.
All weight matrices are initialized by Glorot Initialization (Glorot & Bengio, 2010), and the biases are initialized with 0 .
We use 300D GloVe 6B pre-trained vectors (Pennington et al., 2014) to initialize the word embeddings in x .
The Out-of-Vocabulary words in the training set are randomly initialized by uniform distribution between ( − 0 .
The word embeddings are ﬁne-tuned during the training.
The Dropout (Srivastava et al., 2014) keep probability and the L2 regularization weight decay factor γ are set to 0 .
75 and 5 × 10 − 5 , respectively.
The number of hidden units is 300 .
The unspeciﬁed activation functions in all models are set to Relu (Glorot et al., 2011).
In Table 1, we report the number of parameters, and training/test accuracies of all baselines plus the methods from the ofﬁcial leaderboard.
For fair comparison, we use 480D Bi-BloSAN, which leads to the similar parameter number with that of baseline encoders.
Bi-BloSAN achieves the best test accuracy (similar to DiSAN) among all the sentence encoding models on SNLI.
In particular, compared to the RNN models, Bi-BloSAN outperforms Bi-LSTM encoder, Bi-LSTM with attention and deep gated attention by 2 . 4% , 1 . 5%
and 0 .
2% , respectively.
Bi-BloSAN can even perform better than the semantic tree based models: SPINN-PI encoder ( +2 .
&Tree-based CNN encoder ( +3 .
, and the memory network based model: NSE encoder ( +1 . 1% ).
Additionally, Bi-BloSAN achieves the best performance among the baselines which are based on RNN/CNN/SAN.
It outperforms Bi-LSTM ( +0 . 7% ), Bi-GRU ( +0 . 8% ), Bi-SRU ( +0 . 9% ), multi-CNN ( +2 .
5% ), Hrchy-CNN ( +1 . 8% ) and multi-head attention ( +1 . 5% ).
[Table]
In addition, we compare time cost and memory consumption of all the baselines in Table 2.
Compared to DiSAN with the same test accuracy, Bi-BloSAN is much faster and more memory efﬁcient.
In terms of training and inference time, Bi-BloSAN is 3 ∼ 4 × faster than the RNN models (BiLSTM, Bi-GRU, etc.).
It is as fast as CNNs and multi-head attention but substantially outperforms them in test accuracy.
In terms of training memory, Bi-BloSAN requires similar GPU memory to the RNN-based models and multi-head attention, which is much less than that needed by DiSAN.
Finally, we conduct an ablation study of Bi-BloSAN in Table 3.
In particular, we evaluate the contribution of each part of Bi-BloSAN by the change of test accuracy after removing the part.
The removed part could be: 1) local context representations h , 2) global context representations
E , 3) the context fusion module (mBloSA) or 4) all fundamental modules appeared in this paper.
The results show that both the local and global context representations play signiﬁcant roles in BiBloSAN.
They make Bi-BloSAN surpass the state-of-the-art models.
Moreover, mBloSA improves the test accuracy from 83 .
1% to 85 .
Source2token self-attention performs much better than vanilla attention, and improves the test accuracy by 3 . 3% .
[Table] [Table]
[Table]
R EADING C OMPREHENSION
Given a passage and a corresponding question, the goal of reading comprehension is to ﬁnd the correct answer from the passage for the question.
We use the Stanford Question Answering Dataset (Rajpurkar et al., 2016) (SQuAD) 3 to evaluate all models.
SQuAD consists of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to each question is a segment of text, or a span, from the corresponding passage.
Since Bi-BloSAN and other baselines are designed for sequence encoding, such as sentence embedding, we change the task from predicting the answer span to locating the sentence containing the correct answer.
We build a network structure to test the power of sequence encoding in different models to ﬁnd the correct answers.
The details are given in Appendix B.
Training Setup : We use Adadelta optimizer to minimize the cross-entropy loss plus L2 regularization penalty, with batch size of 32 .
The network parameters and word embeddings initialization methods are same as those for SNLI, except that both the word embedding dimension and the number of hidden units are set to 100 .
We use 0 .
8 dropout keep probability and 10 − 4 L2 regularization weight decay factor.
We evaluate the Bi-BloSAN and the baselines except DiSAN because the memory required by
DiSAN largely exceeds the GPU memory of GTX 1080Ti (11GB).
The number of parameters, per epoch training time and the prediction accuracy on development set are given in Table 4.
[Table]
Compared to RNN/CNN models
[Table]
Bi-BloSAN achieves state-of-the-art prediction accuracy in this modiﬁed task.
Bi-BloSAN shows its competitive context fusion and sequence encoding capability compared to Bi-LSTM, Bi-GRU, Bi-SRU but is much more time-efﬁcient.
In addition, Bi-BloSAN signiﬁcantly outperforms multi-CNN and multi-head attention
[Table]
4. S EMANTIC R ELATEDNESS
The goal of semantic relatedness is to predict the similarity degree of a given pair of sentences.
Unlike the classiﬁcation problems introduced above, predicting the semantic relatedness of sentences is a regression problem.
We use s 1 and s 2 to denote the encodings of the two sentences, and assume that the similarity degree is between [1 , K ] .
Following the method introduced by Tai et al.
(2015), the concatenation of s 1 (cid:12) s 2 and | s 1 − s 2 | is used as the representation of sentence related- ness.
This representation is fed into a 300D fully connected layer, followed by a K-unit output layer with softmax to calculate a probability distribution ˆ p .
The details of this regression problem can be found in Appendix C. We evaluate all models on Sentences Involving Compositional Knowledge (SICK) 4 dataset, where the similarity degree is denoted by a real number in the range of [1 , 5] .
SICK comprises 9,927 sentence pairs with 4,500/500/4,927 instances for training/dev/test sets.
Training Setup :
The optimization objective of this regression problem is the KL-divergence plus the L2 regularization penalty.
We minimize the objective using Adadelta with batch size of 64 .
The network parameters and word embeddings are initialized as in SNLI experiment.
The keep probability of dropout is set to 0 .
7 , and the L2 regularization weight decay factor is set to 10 − 4 .
[Table]
The performances of all models are shown in Table 5, which shows that Bi-BloSAN achieves stateof-the-art prediction quality.
Although Dependency Tree-LSTM and DiSAN obtain the best performance, the Tree-LSTM needs external semantic parsing tree as the recursive input and expensive recursion computation
[Table]
and DiSAN requires much larger memory for self-attention calculation.
By contrast, Bi-BloSAN, as a RNN/CNN-free model, shows appealing advantage in terms of memory and time efﬁciency.
Note that, performance of Bi-BloSAN is still better than some common models, including Bi-LSTM, CNNs and multi-head attention.
4. S ENTENCE C LASSIFICATIONS
The goal of sentence classiﬁcation is to correctly predict the class label of a given sentence in various scenarios.
We evaluate the models on six sentence classiﬁcation benchmarks for various NLP tasks, such as sentiment analysis and question-type classiﬁcation.
They are listed as follows.
• CR 5 : Customer reviews (Hu & Liu, 2004) of various products (cameras etc.).
This task is to predict whether the review is positive or negative.
• MPQA 6 : Opinion polarity detection subtask of the MPQA dataset (Wiebe et al., 2005).
• SUBJ 7 : Subjectivity dataset (Pang & Lee, 2004), which includes a set of sentences.
The corresponding label indicates whether each sentence is subjective or objective.
• TREC 8 : TREC question-type classiﬁcation dataset (Li & Roth, 2002) which coarsely classiﬁes the question sentences into six types.
• SST-1 9 : Stanford Sentiment Treebank (Socher et al., 2013), which is a dataset consisting of movie reviews with ﬁve ﬁne-grained sentiment labels, i.e., very positive, positive, neutral, negative and very negative.
• SST-2 : Stanford Sentiment Treebank (Socher et al., 2013) with binary sentiment labels.
Compared to SST-1, SST-2 removes the neutral instances, and labels the rest with either negative or positive.
Note that only SST-1 and SST-2 have the standard training/dev/test split, and TREC has the training/dev split.
We implement 10-fold cross validation on SUBJ, CR and MPQA because the original datasets do not provide any split.
We do not use the Movie Reviews (Pang & Lee, 2005) dataset because the SST-1/2 are extensions of it.
Training Setup : We use the cross-entropy loss plus L2 regularization penalty as the optimization objective.
We minimize it by Adam with training batch size of 32 (except DiSAN, which uses batch size of 16 due to the limit of GPU memory).
The network parameters and word embeddings are initialized as in SNLI experiment.
To avoid overﬁtting on small datasets, we decrease the dropout keep probability and the L2 regularization weight decay factor γ to 0 .
6 and 10 − 4 , respectively.
[Table]
The prediction accuracies of different models on the six benchmark datasets are given in Table 6.
Bi-BloSAN achieves the best prediction accuracies on CR, MPQA and TREC, and state-of-theart performances on SUBJ, SST-1 and SST-2 datasets (slightly worse than the best performances).
Although Bi-BloSAN performs a little bit worse than the RNN models on SUBJ and SST-1, it is much more time-efﬁcient than them.
Additionally, on the SST-2 dataset, Bi-BloSAN performs slightly worse than DiSAN in terms of prediction accuracy ( − 0 .
but obtains a signiﬁcantly higher memory utility rate.
We visualize the progress of training models on CR dataset in Figure 5.
The convergence speed of
Bi-BloSAN is ∼ 6 × and ∼ 2 × faster than Bi-LSTM and DiSAN respectively.
Although Bi-BloSAN is less time-efﬁcient than CNN and multi-head attention, it has much better prediction quality.
4. A NALYSES OF T IME C OST AND M EMORY C ONSUMPTION
To compare the efﬁciency-memory trade-off for each model on sequences of different lengths, we generate random tensor data, and feed them into the different sequence encoding models.
The models we evaluate include Bi-LSTM, Bi-GRU, Bi-SRU, CNN, multi-head attention, DiSAN and
Bi-BloSAN.
The shape of the random data is [ batch size , sequence length , features number ].
We ﬁx the batch size to 64 and the features number to 300 , then change the sequence length from 16 towith a step size 16 .
We ﬁrst discuss the time cost vs. the sequence length.
As shown in Figure 6(a), the inference time of
Bi-BloSAN is similar to those of multi-head attention and multi-CNN, but Bi-BloSAN outperforms
[Figure]
both by a large margin on prediction quality in previous experiments.
Moreover, Bi-BloSAN is much faster than the RNN models (Bi-LSTM, Bi-GRU, BI-SRU).
In addition, although DiSAN requires less training time than the RNN models in the experiments above, it is much slower during the inference phase because the large memory allocation consumes a great amount of time.
By contrast, the block structure of Bi-BloSAN signiﬁcantly reduces the inference time.
The GPU memory consumption vs. the sequence length for each model is visualized in Figure 6(b).
DiSAN is not scalable because its memory grows explosively with the sequence length.
Bi-BloSAN is more memory-efﬁcient and scalable than DiSAN as the growth of its memory is nearly linear.
Although Bi-BloSAN consumes more memory than the RNN models, it experimentally has better time efﬁciency and prediction quality.
Since multi-head attention uses multiplicative attention, it requires less memory than all additive attention based models, such as DiSAN and Bi-BloSAN, but multiplicative attention based models usually perform worse than additive attention based models.
C ONCLUSIONS
This paper presents an attention network, called bi-directional block self-attention network (BiBloSAN), for fast, memory-efﬁcient and RNN/CNN-free sequence modeling.
To overcome large memory consumption of existing self-attention networks, Bi-BloSAN splits the sequence into several blocks and employs intra-block and inter-block self-attentions to capture both local and long- range context dependencies, respectively.
To encode temporal order information, Bi-BloSAN applies forward and backward masks to the alignment scores between tokens for asymmetric selfattentions.
Our experiments on nine benchmark datasets for various different NLP tasks show that Bi-BloSAN can achieve the best or state-of-the-art performance with better efﬁciency-memory trade-off than existing RNN/CNN/SAN models.
Bi-BloSAN is much more time-efﬁcient than the RNN models (e g.,
Bi-LSTM, Bi-GRU, etc.), requires much less memory than DiSAN, and signiﬁcantly outperforms the CNN models and multi-head attention on prediction quality.  A CKNOWLEDGMENTS
This research was funded by the Australian Government through the Australian Research Council (ARC) under grants 1) LP160100630 partnership with Australia Government Department of
Health and 2) LP150100671 partnership with Australia Research Alliance for Children and Youth (ARACY) and Global Business College Australia (GBCA).
We also acknowledge the support of
NVIDIA Corporation with the donation of GPU used for this research
A T HE S ELECTION OF B LOCK L ENGTH
In mBloSA the length r of each block is a hyper-parameter that determines memory consumption of mBloSA.
To minimize the memory consumption, we propose an approach that calculates the optimized block length r as follows.
We ﬁrst introduce the method for determining r for a dataset that has ﬁxed sentence length n .
Given the sentence length n and the block number m = n/r , we have the following facts: 1) the major memory consumption in mBloSA is dominated by the masked self-attentions g m ( · , M ) ; 2) the memory consumption of the masked self-attention is proportional to the square of the sentence length; and 3) mBloSA contains m masked self-attention with a sequence length of r , and 1 masked self-attention with a sequence length of m .
Therefore, the memory ξ required by mBloSA can be calculated by
[Equation]
By setting the gradient of ξ w.r.t.
r to zero, we know that the memory consumption ξ is minimum when r = 3 √ 2 n .
Second, we propose a method for selecting r given a dataset with the sentence lengths that follow a normal distribution N ( µ, σ 2 ) .
We consider the case where mini-batch SGD with a batch size of B or its variant is used for training.
We need to calculate the upper bound of the expectation of the maximal sentence length for each mini-batch.
Let us ﬁrst consider B random variables [ X 1 , X 2 , . . .
, X B ] in the distribution N (0 , σ 2 ) .
The goal is to ﬁnd the upper bound of the expectation of random variable Z + µ , where Z is deﬁned as
[Equation]
By Jensen’s inequality,
[Equation]
Eq.(21) leads to
[Equation]
Let t = √ 2ln B σ
and we obtain the following upper bound.
[Equation]
[Equation]
Hence, the upper bound of the expectation of the maximal sentence length among all the B sentences in each mini-batch is σ √ 2 ln B + µ .
Therefore, the block length r is computed by
[Equation]
B N ETWORK S ETUP FOR M ACHINE C OMPREHENSION
Each sample in the Stanford Question Answering Dataset (SQuAD) (Rajpurkar et al., 2016) is composed of three parts, i.e., a passage consisting of multiple sentences, a question sentence and a span in the passage indicating the position of the answer.
In order to evaluate the performance of sentence embedding models, we change the task from predicting the span of the answer to ﬁnding the sentence containing the correct answer.
Given a passage consisting of m sentences [ s 1 , s 2 , . . .
, s m ] where s k =
[ x k 1 , x k 2 , . .
, x kn ] , and the embedded question token sequence q = [ q 1 , q 2 , . . .
, q l ] , the goal is to predict which sentence in the m sentences contains the correct answer to the question q .
The neural net we use to evaluate different sequence encoding models is given in Figure 7.
First, we process each sentence from the passage by a context fusion layer with shared parameters, followed by a source2token self-attention with shared parameters, which outputs a vector representation of the sentence.
Therefore, the m sentences are represented by m vectors [ u 1 , u 2 , . . .
, u m ] .
The question sentence q is compressed into a vector representation q using source2token self-attention.
Second, we combine each sentence u k with q by concatenating u k , q , u k − q and u k (cid:12) q , i.e.,
[Equation]
Then c =
[ c 1 , c 2 , . .
, c m ] is fed into another context fusion layer that explores sentence-level dependencies.
Finally, the resultant output representation of each sentence is separately fed into a fully connected layer to compute a scalar score indicating the possibility of the sentence containing the answer.
A softmax function is applied to the scores of all m sentences, to generate a probability distribution ˆ p ∈ R m for cross-entropy loss function.
The sentence with the largest probability is predicted as the sentence containing the answer.
C L OSS OF R EGRESSION P ROBLEM
Following the setting introduced by Tai et al.
and given a predicted probability distribution ˆ p as the output of a feedforward network, the regression model predicts the similarity degree as
[Equation]
(25) where β
= [1 , 2 , . .
, K ] .
The ground-truth similarity degree y should be mapped to a probability distribution p =
[ p i ]
Ki =1 as the training target, where p needs to fulﬁll y =
β T p .
The mapping can be deﬁned as
[Equation]
otherwise , i = 1 , 2 , . .
K.
We use KL-divergence between p and ˆ p as our loss function, i.e.,
[Equation]
(27) where the p ( k ) and ˆ p ( k ) represent the target and predicted probability distributions of the k -th sample, respectively.
D R ELATED W ORKS
Recently, several structured attention mechanisms (Kim et al., 2017; Kokkinos & Potamianos, 2017) are proposed for capturing structural information from input sequence(s).
When applied to selfattention, structured attentions share a similar idea to self-alignment attention (Hu et al., 2017) and multi-head attention (Vaswani et al., 2017) with one head, which aims to model the dependencies between the tokens.
Similar to the attention from multiple perspectives in multi-head attention, multi-perspective context matching (Wang et al., 2016) explores the dependencies between passage and question from multiple perspectives for reading comprehension, while self-attentive structure (Lin et al., 2017) embeds sentences from various perspectives to produce matrix representations of the sentences.
In recursive models, self-attention over children nodes (Teng & Zhang, 2017) can provide effective input for their parent node that has no standard input (Tai et al., 2015) as long as it is a non-leaf node in a semantic constituency parsing tree.
Bi-BloSA and hierarchical attention network (Yang et al., 2016) have similar structure, i.e., both stack two-layer attention mechanisms from bottom to top.
However, they are different in three respects: 1) Bi-BloSA aims to learn context-aware representation for each token, while hierarchical attention is designed for document embedding; 2) the input of Bi-BloSA is a sentence, while the intput of hierarchical attention is a document composed of multiple sentences; and 3) the hierarchical attention performs vanilla attention twice, i.e., token-level attention on each sentence and sentence-level attention.
Bi-BloSA, however, applies masked self-attention twice, i.e., intra-block self-attention and inter-block self-attention.
Additionally, Bi-BloSA uses a feature fusion gate for each token to combine local and global context, and positional masks to encode temporal order information.