P O N ET : P OOLING N ETWORK FOR E FFICIENT T OKEN M IXING IN L ONG S EQUENCES
Anonymous authors
Paper under double-blind review
A BSTRACT
Transformer-based models have achieved great success in various NLP, vision, and speech tasks.
However, the core of Transformer, the self-attention mechanism, has a quadratic time and memory complexity with respect to the sequence length, which hinders applications of Transformer-based models to long sequences.
Many approaches have been proposed to mitigate this problem, such as sparse attention mechanisms, low-rank matrix approximations and scalable kernels, and token mixing alternatives to self-attention.
We propose a novel Pooling
Network ( PoNet ) for token mixing in long sequences with linear complexity.
We design multi-granularity pooling and pooling fusion to capture different levels of contextual information and combine their interactions with tokens.
On the
Long Range Arena benchmark, PoNet signiﬁcantly outperforms Transformer and achieves competitive accuracy, while being only slightly slower than the fastest model, FNet, across all sequence lengths measured on GPUs.
We also conduct systematic studies on the transfer learning capability of PoNet and observe that
PoNet achieves 95.7% of the accuracy of BERT on the GLUE benchmark, outperforming FNet by 4.5% relative.
Comprehensive ablation analysis demonstrates effectiveness of the designed multi-granularity pooling and pooling fusion for token mixing in long sequences and efﬁcacy of the designed pre-training tasks for PoNet to learn transferable contextualized language representations 1 .
I NTRODUCTION
Transformer (Vaswani et al., 2017) has become the state-of-the-art (SOTA) architecture for sequence modeling in a wide variety of ﬁelds, including natural language processing (NLP), computer vision, speech processing, applications to genomics data, etc.
The key reason for Transformer’s success is its self-attention mechanism, which computes dot-product between input representations for each pair of positions in the full sequence.
Proved to be greatly effective in learning contextualized representations, Transformer becomes the backbone for dominant pre-trained language models (PLM) in NLP, such as BERT (Devlin et al., 2019) and RoBERTa (Liu et al., 2019).
These
PLMs demonstrate strong transfer learning capabilities and have achieved SOTA widely on NLP tasks.
However, self-attention has quadratic time and memory complexity to the input sequence length (Vaswani et al., 2017), which becomes the bottleneck for applying the vanilla Transformer to long sequence modeling tasks and for scaling up Transformer-based PLMs.
A broad spectrum of approaches have been proposed to address the efﬁciency problem of selfattention, as summarized in (Tay et al., 2020).
One major direction approximates the dense full self-attention using techniques such as introducing sparsity (Child et al., 2019; Beltagy et al., 2020;
Zaheer et al., 2020; Zhang et al., 2021a;
Yang et al., 2021; Zhu & Soricut, 2021; Zhang et al., 2021b), low-rank approximations of the softmax attention matrix (Katharopoulos et al., 2020;
Wang et al., 2020a; Zhu & Soricut, 2021; Xiong et al., 2021; Peng et al., 2021; Choromanski et al., 2021), and locality sensitive hashing (Kitaev et al., 2020).
These approximation approaches exploit observations that token interactions have strong locality, hence the importance and in turn the attention should decrease with the increase of the distance between a query token and a key token.
Several of these works achieve O ( N ) theoretical complexity.
However, these models often require selecting relatively large local regions for ﬁne-grained attention in order to approximate full self-attention, hence the scaling constants hidden by O ( N ) are often large and hinder signiﬁcant improvements in speed and memory usage.
It is observed that performance of the approximation approaches is usually inversely related to their speed (Beltagy et al., 2020; Zaheer et al., 2020).
Another major direction replaces the self-attention structure with more efﬁcient structures, such as
MLP-Mixer (Tolstikhin et al., 2021), FNet (Lee-Thorp
et al., 2021), AFT (Zhai et al., 2021), and
Fastformer (Wu et al., 2021).
Despite signiﬁcant accelerations gained by efﬁcient transformers, they are rarely evaluated both on effectiveness of the inductive bias and the transfer learning capability, except BigBird, FNet, and Nystr¨omformer 2 as we understand.
BigBird-Base (Zaheer et al., 2020) outperforms BERT-Base on the GLUE benchmark (Wang et al., 2019) without signiﬁcant speedup (slowdown for input length < 3 K ) (Tay et al., 2021).
FNet-Base trains 80% faster than BERT on
GPUs at 512 input lengths but only achieves 92% of BERT-Base’s accuracy (Lee-Thorp et al., 2021).
In this work, we propose a novel Pooling Network ( PoNet ), aiming to simultaneously advance long sequence modeling capacity and transfer learning capabilities while improving speed and memory efﬁciency.
PoNet replaces the O ( N 2 ) complexity self-attention with a O ( N ) complexity multi-granularity pooling block.
We design multi-granularity pooling and pooling fusion to model different levels of token interactions.
Multi-granularity pooling incorporates three types of pooling, from coarse-grained to ﬁne-grained, in each sublayer.
Global aggregation aggregates information of the entire sequence into a single token.
Segment max-pooling captures the paragraph or sentence level information.
Local max-pooling captures the more important local information.
These three poolings are fused to produce the output feature of the multi-granularity pooling block.
Then through the residual connection, this output feature is further aggregated into each token.
The contributions of this paper are summarized as follows:
• We propose a novel PoNet architecture as a drop-in replacement for self-attention in Transformer, achieving linear time and memory complexity.
We propose multi-granularity pooling and pooling fusion to capture different levels of contextual information and comprehensively model token interactions.
To the best of our knowledge, our work is the ﬁrst to explore the full potential of the simple pooling mechanism for token mixing and modeling long-range dependencies.
• Extensive evaluations show that PoNet achieves competitive performance on the Long Range Arena benchmark (Tay et al., 2021) and signiﬁcantly outperforms Transformer by +2.28 absolute (+3.9% relative) on accuracy, with efﬁciency up to 9 times faster and 10 times smaller than Transformer on GPU.
Also, PoNet demonstrates competitive transfer learning capabilities, with PoNet-Base reaching 95.7% of the accuracy of BERT-Base on the GLUE benchmark.
Ablation analysis further proves effectiveness of designed multi-granularity pooling and pre-training tasks.
R ELATED W ORK
Efﬁcient Transformer Variants Among the models to approximate full self-attention, Longformer (Beltagy et al., 2020)
( O ( N ) ) sparsiﬁes the full self-attention into three attention patterns of sliding window, dilated sliding window, and global attention.
BigBird (Zaheer et al., 2020) ( O ( N ) ) combines global attention, local attention, and random attention.
Poolingformer (Zhang et al., 2021a)
( O ( N ) ) uses a two-level attention schema, with the ﬁrst level using a smaller sliding window to aggregate local information and the second level using a larger window with pooling attention to reduce time and memory cost.
Focal Transformer (Yang et al., 2021)
( O ( N ) ) uses both ﬁne-grained local interactions and coarse-grained global interactions to balance efﬁciency and effectiveness of capturing short- and long-range dependencies.
Transformer-LS (Zhu et al., 2021)
( O ( N ) ) approximates the full attention by aggregating long-range attention via dynamic projections and short-term attention via segment-wise sliding window.
H-Transformer-1D (Zhu & Soricut, 2021)
( O ( N ) ) exploits a matrix structure similar to Hierarchical Matrix.
AdaMRA (Zhang et al., 2021b)
( O ( N ) ) leverages a multi-resolution multi-head attention mechanism and kernel attention.
Luna (Ma et al., 2021)
( O ( N ) ) introduces an additional ﬁxed length sequence served as query to attend to the original input while the output is served as key and value to attend to the original input.
Apart from the sparse attention models, other approximation approaches explore locality sensitive hashing and matrix approximation methods.
Reformer (Kitaev et al., 2020)
O ( NlogN ) replaces self-attention with locality sensitive hashing.
Performer (Choromanski et al., 2020; 2021)
( O ( N ) ) approximates softmax attention by leveraging random features.
Linformer (Wang et al., 2020a) approximates the self-attention matrix with a low-rank factorization.
Nystr¨omformer
(Xiong et al., 2021)
( O ( N ) ) approximates the softmax attention with the Nystr¨om method by sampling a subset of columns and rows.
However, these approximation methods have strengths on certain tasks and may cause accuracy degradation on many other tasks.
Our work is in another line of research on replacing self-attention with more efﬁcient token mixing mechanisms.
MLP-Mixer (Tolstikhin et al., 2021)
( O ( N 2 ) ) applies two separate linear transformations on the hidden state dimension and the sequence dimension.
FNet (Lee-Thorp et al., 2021)
( O ( NlogN ) ) replaces the self-attention sublayer with 2D-FFT mixing sublayer.
AFTlocal/conv (Zhai et al., 2021)
( O ( sN ) , s < N ) ﬁrst combines the key and value with a set of learned position biases and then combines the query with this result via element-wise multiplication.
Fastformer (Wu et al., 2021)
( O ( N ) )
ﬁrst models global context via additive attention then models interactions between global context and input representations through element-wise product.
Pre-training Tasks It has been observed that both underlying model architecture and pre-training are crucial to performance of PLMs.
BERT (Devlin et al., 2019) with a Transformer encoder is pretrained with masked language modeling (MLM) and next sentence prediction (NSP) tasks on largescale unlabeled text corpora including the English Wikipedia and BooksCorpus.
MLM predicts the masked token from context.
NSP predicts whether a sentence pair is contiguous or not in the original source.
Many approaches are proposed to improve these two tasks and show that more challenging pre-training tasks may help PLMs learn better and more transferable language representations.
Whole word masking (WWM) (Devlin et al., 2019; Cui et al., 2019) and SpanBERT (Joshi et al., 2020) outperform BERT on many tasks.
WWM simultaneously masks all WordPiece tokens belonging to the same word and forces the model to recover a complete whole word.
SpanBERT randomly samples contiguous spans inside of individual tokens and augments MLM with a new task to predict the entire masked span.
RoBERTa
(Liu et al., 2019) reports ineffectiveness of NSP and removes it from pre-training.
ALBERT (Lan et al., 2020) replaces NSP with a sentence-order prediction (SOP) task to predict whether two consecutive sentences are in the right order or not, for learning ﬁne-grained inter-sentence coherence.
StructBERT (Wang et al., 2020b) extends SOP to a new sentence structural objective (SSO) as a ternary classiﬁcation on two sentences ( S 1 , S 2 ) to decide whether S 1 precedes or follows S 2 or the two sentences are noncontiguous.
More challenging tasks for learning inter-sentence relations and document/discourse structures (Iter et al., 2020; Lee et al., 2020;
Ding et al., 2021) show promising performance improvements on PLMs.
M ODEL
Our work is inspired by the External Attention (EA) approach proposed in (Guo et al., 2021) for visual tasks.
An input sequence of tokens x = { x 1 , ..., x N } are mapped to an embedding matrix denoted by H ( ∈ R N × d ) = {
h 1 , ..., h N } , where N is the sequence length and d is the hidden dimension.
EA uses two linear layers to implement external and shared memories, which facilitates learning correlations across all samples and hence serves strong regularization to and improves generalization of the attention mechanism with linear complexity.
We simplify
EA into multi-layer perceptron and softmax , and observe that by infusing the sequence-level information into each token through the denominator term (cid:80)
Nn =1 e h n , softmax provides context modeling capabilities.
However, softmax involves calculations of exponents, which is still slow.
Consequently, we consider using pooling as an alternative to capture contextual information with signiﬁcantly reduced complexity.
We propose a Pooling Network ( PoNet ) as a drop-in replacement for the self-attention sublayer in Transformer, as shown in Figure 1.
PoNet models different levels of contextual information through a multi-granularity pooling (MP) block consisting of three components, namely, global aggregation (GA), segment max-pooling (SMP), and local maxpooling (LMP).
These pooling features are then aggregated through pooling fusion.
3. M ULTI - GRANULARITY P OOLING
In order to capture contextual information at different levels, we design multi-granularity pooling.
First, different linear projections are applied on the input for poolings at different granularities:
[Equation]
[Figure]
where ∗ represents Q g , K g , V g for GA, s for SMP, l
LMP and o fusion in six
W ∗ .
W ∗ ∈ R
d × d and b ∗ ∈ R
d are parameters to be learned.
The different H ∗ are then used for different poolings.
3.1. G LOBAL A GGREGATION
A Global Aggregation (GA) module is carefully designed aiming to both capture the most important global information for each token and also to guarantee an overall linear computational complexity.
We calculate the ﬁrst stage value g for GA by averaging at the sequence level:
[Equation]
Note that g is only a rough representation of the sequence.
Beltagy et al.
Zaheer et al.
(2020) introduced a global attention mechanism, which adds a set of global tokens with randomly initialized values to always attend to the whole sequence.
Inspired by these works, we perform cross-attention on the ﬁrst stage value for GA.
The ﬁrst stage value g is used to perform a query on the input sequence to compute the second stage value g (cid:48) for GA, as follows:
[Equation]
The cross-attention on g enables each token to attend to the whole sequence and hence the resulting second stage value g (cid:48) provides a more accurate sequence representation, compared to g .
Note that since g is a single token with the length equaling 1 , the computational complexity of attention in
Eq. 3 is O ( N ) .
Theoretically, using average-pooling for generating g , the rough representation of the sequence, could keep more information for the next step cross-attention and generating the second stage value for GA, instead of discarding most information and only focusing on the most salient features as the max-pooling function does.
This hypothesis is veriﬁed as we observe a better model performance from using average-pooling for generating g than max-pooling.
In contrast, the outputs of SMP and LMP are directly used as ﬁnal representations, hence we choose max-pooling for SMP and LMP, which is also empirically veriﬁed to produce a better model performance.
3.1. S EGMENT M AX - POOLING
The information loss from compressing a long sequence into a single global token could be enormous and hence become detrimental to the sequence modeling capacity.
We introduce an intermediate level between tokens and the global token by segmenting an input sequence into segments, exploring prior knowledge of structure in the data when available, and introduce Segment
Max-pooling (SMP).
The use of structural information for segmentation is adjustable for different tasks, which is described in detail in Section 4.
As explained earlier, max-pooling is performed on each segment at each dimension j ∈ [1 , d ] , as follows, where K denotes the number of segments:
[Equation]
3.1. L OCAL M AX - POOLING
Many prior works (Zhu & Soricut, 2021; Yang et al., 2021) demonstrate the crucial role of capturing local information for the sequence modeling capabilities.
We introduce Local Max-pooling (LMP), designed as a standard max-pooling over sliding windows, to capture contextual information from neighboring tokens for each token.
Different from GA and SMP, the window for LMP is overlapped.
Similar to GA, LMP is also applied at the sequence level and the left and right boundaries of the input sequence are padded to ensure that the output length equals the original input length.
The LMP values L ∈ R N × d are computed.
The size and stride of the sliding window are set to 3 and 1 in all our experiments unless stated otherwise.
3.1. P OOLING F USION
First, we model the interactions between GA and each token by computing G n through the elementwise product between the second stage value of GA g (cid:48) and each token, as follows:
[Equation]
The reason for conducting Eq.
7 instead of directly using g (cid:48) as the output is to avoid all tokens from sharing the same global token.
Otherwise, it will cause the token representations to converge to the same global token in the subsequent addition fusion layer.
This effect, in turn, will make the token representations become more homogeneous, and consequently degrade performance on tasks such as sentence-pair classiﬁcations.
The SMP token S is shared by all tokens in each segment.
Hence, for the same rationale of mixing the global token with each token, the same operation is conducted to mix the SMP token with each token and to compute S (cid:48) as:
[Equation]
where k ( n ) denotes the segment index of the n -th token.
The above three features are added up as the ﬁnal output of our multi-granularity pooling block, as illustrated in Figure 1:
[Equation]
where P is used to replace the original self-attention output of Transformer.
We compare PoNet to related models, including Fastformer (Wu et al., 2021), Luna (Ma et al., 2021), Longformer (Beltagy et al., 2020) and BigBird (Zaheer et al., 2020), in detail in Appendix D.
3. C OMPLEXITY A NALYSIS
We only analyze the computational complexity of the proposed multi-granularity pooling block, since we only replace the self-attention sublayer in Transformer with this block and keep other modules in Transformer unchanged.
The six HW ∗ in Eq.
1, where ∗ represents Q g , K g , V g for
GA, s for SMP, l for LMP, and o for pooling fusion, require 6 Nd 2 computations.
GA requires 2 Nd ops (Eq. 3), SMP and LMP have no matrix multiplication, and Pooling Fusion requires 2 Nd ops (Eq. 7 and Eq.
Hence, the total number of multiplication ops is 6 Nd 2 + 4 Nd .
We further simplify computations.
By switching the order of Eq. 1 and Eq. 2 into ﬁrst performing the average pooling then the afﬁne function, computation can be reduced from Nd 2 to d 2 .
Adding g (cid:48) and S ﬁrst and then performing element-wise product can reduce 2 Nd to Nd , compared to conducting
Eq.
7 and Eq.
8 separately.
After the simpliﬁcations, the total number of multiplication ops is (5 N + 1) d 2 + 3 Nd .
The multi-granularity pooling block hence has linear time and memory complexity with respect to the input length.
[Table]
., 2021), with (3) are from (Xiong et al., 2021).
We implement our PoNet and re-implement FNet based on the PyTorch codebase from (Xiong et al., 2021) and use the same experimental conﬁgurations to ensure a fair comparison.
For each group, the best result for each task and AVG are bold-faced.
E XPERIMENTS
We ﬁrst evaluate PoNet on the Long Range Arena (LRA) benchmark (Tay et al., 2021) and compare
PoNet to the vanilla Transformer and a series of efﬁcient transformer variants on accuracy, training speed, and memory usage.
Next, we study the transfer learning capability of PoNet in the commonly used paradigm of pre-training followed by ﬁne-tuning.
We evaluate the ﬁne-tuning performance on the GLUE benchmark (Wang et al., 2019) as well as a set of long-text classiﬁcation tasks.
All baseline models and PoNet use the same “Base” model conﬁguration as BERT-Base (Devlin et al., 2019).
PoNet-Base has 124M parameters (see the ﬁrst paragraph in Appendix A for more details).
More experimental details and results are in Appendices.
4. L ONG -R ANGE A RENA B ENCHMARK
Comparison on Accuracy The LRA benchmark is designed to assess the general capabilities of capturing long-range dependencies.
LRA consists of six tasks spanning structural reasoning ( ListOps ), similarity reasoning (Byte-level Text classiﬁcation and document Retrieval , Image classiﬁcation), and visual-spatial reasoning ( Pathﬁnder ).
We use the PyTorch codebase from (Xiong et al., 2021) 3 to implement FNet and our PoNet and evaluate on LRA after ﬁrst replicating the results from Nystr ¨ omformer in (Xiong et al., 2021).
We keep the same hyperparameter setting as used by (Xiong et al., 2021) for all of our LRA evaluations and report results on ﬁve tasks in
Table 1 4 .
All the baseline models are summarized in Section 2 and the Linear variant (Lee-Thorp et al., 2021)
denotes replacing the self-attention sublayer with two linear projections, one applied to the hidden dimension and one applied to the sequence dimension.
Note that due to different code implementations, results for same models could differ across groups in Table 1 (see Appendix A. for details).
It is important to point out that all models in the third group are implemented with the same codebase from (Xiong et al., 2021) and the same experimental conﬁgurations to ensure a fair comparison within this group.
As can be seen in Table 1, compared to the ﬁrst and second groups and the cited results in the third group marked with (3), PoNet achieves competitive performance on
LRA.
PoNet outperforms the vanilla Transformer by +2.28 (61.05 over 58.77) and Nystr¨omformer by +2.10 on the average accuracy and consistently accomplishes better performance on all tasks compared to Transformer, Performer, Reformer, Linformer, FNet, and Nystr ¨ omformer, except slightly weaker than Transformer on the Pathﬁnder task.
It is reasonable to conclude that PoNet outperforms BigBird on LRA since the margin +2.28 from PoNet over Transformer in the third
[Table]
group is signiﬁcantly larger than the margin +0.62 from BigBird over Transformer in the ﬁrst group.
To the best of our knowledge, PoNet achieves the third best accuracy on LRA against Transformer and recent efﬁcient transformers, only lower than 63.09 from AdaMRA (Zhang et al., 2021b) 5 and 61.41 from H-Transformer-1D (Zhu & Soricut, 2021).
Comparison on Speed and Memory Consumption Table 2 compares the GPU training speed and peak memory consumption of PoNet to Transformer, Performer, Nystr¨omformer, and FNet on a single NVIDIA V100 chip, on input sequence lengths from 512 up to 16384.
We observe that
PoNet is the second fastest model and consumes the second smallest memory footprint in the group, consistently on all sequence lengths, much faster than Transformer, Performer, and Nystr ¨ omformer and lighter than them, and only slightly slower and heavier than FNet.
Also, the speedup from PoNet over Transformer escalates on longer input sequence lengths.
4. T RANSFER L EARNING
The paradigm of pre-training followed by ﬁne-tuning has been extensively applied and accomplished
SOTA results in a wide variety of NLP tasks.
Therefore, it is critical to evaluate the transferability of
PoNet.
We perform pre-training on PoNet and evaluate the ﬁne-tuning performance on the GLUE benchmark and a set of long-text classiﬁcation benchmarks.
To facilitate a fair comparison on the transfer learning capability, we pre-train BERT, FNet, and PoNet with the same MLM (Devlin et al., 2019) and sentence structural objective (SSO) as in StructBERT (Wang et al., 2020b) on the same English Wikitext-103 (100M words) and
BooksCorpus (800M words) datasets 6 .
The total pre-training loss is L =
L MLM + L SSO .
All three models are base-uncased and pre-trained using the same conﬁguration (Appendix A.2).
For PoNet, we use the natural paragraph segmentations in the datasets to compute SMP.
Figure 2 illustrates validation accuracy of the same MLM and SSO tasks from BERT, FNet, and PoNet pre-training.
MLM accuracy from PoNet is only slightly worse than that from BERT while the gap on SSO accuracy between them is a bit larger.
PoNet achieves signiﬁcantly better MLM and SSO accuracy than FNet, consistent with its better sequence modeling ability shown on LRA.
Results on GLUE
The GLUE benchmark covers a diverse range of challenging natural language understanding tasks and is widely adopted for evaluating transfer learning models.
The tasks can be split into two groups, single-sentence tasks including CoLA and SST-2, and sentence-pair tasks
[Figure] [Figure]
[Table]
including MRPC, QQP, STS-B, MNLI, QNLI
[Table]
and RTE 7 .
The special token “[SEP]” is used as the segment separator.
For ﬁne-tuning PoNet on GLUE, inputs to single-sentence tasks include three segments as “[CLS]” Sentence-1 “[SEP]”; whereas inputs to sentence-pair tasks include ﬁve segments “[CLS]” Sentence-1 “[SEP]”
Sentence-
“[SEP]”.
These segments are used for computing
SMP (Section 3).
Table 3 shows the results for the best base learning rate (no early stopping) on the GLUE Validation split (see Appendix A.3 for more details), providing a fair comparison since all three models are pre-trained and ﬁne-tuned with the same pre-training data/tasks/hyperparameters with 340K steps.
Table 3 shows that PoNet achieves 76.80 AVG score, reaching 95.7% of the accuracy of BERT on GLUE (80.21) and outperforming FNet by 4.5% relatively.
These performance comparisons are consistent with the pre-training accuracies shown in Figure 2.
The results also prove that PoNet is equally competitive in both single-sentence and sentence-pair tasks.
Results on Long-text Classiﬁcation We also evaluate the ﬁne-tuning performance of the pretrained PoNet on four long-text classiﬁcation datasets, including Hyperpartisan news detection (HND) (Kiesel et al., 2019) 8 , IMDb (Maas et al., 2011), Yelp-5 (Zhang et al., 2015), and
Arxiv-11 (He et al., 2019).
As can be seen from Table 4, PoNet-Base outperforms BERT-Base on HND (+8.2 F 1 ) and Arxiv (+0.75 F 1 ) and reaches 99% of BERT-Base’s F 1 on IMDb and Yelp-5.
A BLATION A NALYSIS
We conduct ablation experiments to understand contributions from multi-granularity pooling and pre-training tasks.
By applying leave-one-out on different components of PoNet, we create variants as PoNet w/o Second Stage GA (SS-GA), PoNet w/o GA, PoNet w/o SMP, PoNet w/o LMP, and
PoNet w/o (SMP&LMP).
We also pre-train PoNet with two weaker losses, as L MN =
L MLM + L NSP and Ls OM =
L MLM , where L NSP is the standard NSP loss in BERT pre-training.
We report the pretraining task validation accuracy and GLUE validation split scores on the single-sentence CoLA task
[Table]
[Table]
and the sentence-pair STS-B task in Table 5.
Note that [CLS] from PoNet w/o GA cannot capture the information of the whole sequence, as using [CLS] for classiﬁcation fails on SST, CoLA, and
STS-B tasks, we use max-pooling of the sequence instead for classiﬁcation.
Removing GA (PoNet w/o SS-GA, PoNet w/o GA) signiﬁcantly degrades accuracy on SST and
STS-B, showing that sentence-pair tasks heavily rely on the global information.
MLM accuracy degrades only slightly from PoNet w/o SS-GA but more signiﬁcantly from PoNet w/o GA, indicating that MLM also depends on the global information, but when the rough global information is available (PoNet w/o SS-GA), SMP can compensate sufﬁciently.
Interestingly, PoNet w/o SSGA improves CoLA accuracy.
Removing GA strengthens SMP and LMP learning thus yields better CoLA accuracy since CoLA relies on SMP and LMP much more than GA.
Different from the conclusions on BERT (Devlin et al., 2019; Liu et al., 2019), we ﬁnd ﬁne-tuning performance of PoNet on sentence-pair tasks highly relies on the sentence structural tasks in pre-training.
Weakening SST loss ( L MN , L OM ) weakens GA representation learning while strengthening SMP and LMP learning, producing a signiﬁcant degradation in STS-B accuracy but a signiﬁcant gain in
CoLA accuracy.
Similarly, removing SMP or LMP enhances GA representation learning and hence improves STS-B accuracy while degrading CoLA accuracy.
PoNet w/o SMP&LMP (i.e. GA only) shows a drastic degradation on MLM accuracy and CoLA accuracy (45.36 to 11.36).
These results conﬁrm that all three poolings are important for the modeling capabilities of PoNet.
C ONCLUSION
We propose a novel Pooling Network (PoNet) to replace self-attention with a multi-granularity pooling block, which captures different levels of contextual information and combines them for a comprehensive modeling of token interactions.
Extensive evaluations demonstrate that PoNet achieves both competitive long-range dependency modeling capacity and strong transfer learning capabilities, with linear time and memory complexity.
Future work includes further optimization of model structure and pre-training as well as applying PoNet to a broader range of tasks including generation tasks.
R EPRODUCIBILITY S TATEMENT
All data used in the experiments in this paper are open source.
Readers can refer to the original papers for details of the datasets, which are cited in our paper.
Information on data access can be found in Appendix A. Experimental details are also described in Appendix A. We will release the source code upon the publication of the paper.
Appendices A E XPERIMENT DETAILS
Model Size and How to Reduce It For all experiments for PoNet in this paper, parameters W K g , b K g and W V g , b V g in Equation 1 are shared to reduce the calculations and we observe no performance degradation.
After this sharing of W K g , b K g and W V g , b V g in Equation 1, PoNetBase has 124M parameters.
This is comparable to the 110M parameters for BERT-Base but still a bit larger in the number of parameters.
Note that after this sharing of W K g , b K g and W V g , b V g , our current PoNet model has a total of 5 linear transformations from Equation 1 in the token mixing layer.
Since W Q g will only be performed on a single token (global token) instead of on the whole input sequence, its computational cost can be ignored.
The main computational cost comes from the remaining 4 transformations.
To further reduce the number of parameters of PoNet and reduce computations, we plan to further remove W o since during pooling fusion, the original input sequence could be mixed directly with GA and SMP without transformations. This leaves us with only 3 sets of linear transformation parameters that require the computations of the whole sequence, namely, shared W K g and W V g , W s , and W l . Then PoNet will have the same number of linear transformations as BERT, which has a total of 3 linear transformations in the self-attention layer. A. L ONG -R ANGE A RENA B ENCHMARK E XPERIMENTAL D ETAILS
Implementations and Hyperparameters We use the Pytorch codebase from Xiong et al.
(2021)to implement our PoNet and re-implement FNet and conduct all LRA evaluations to facilitate a fair comparison.
We use exactly the same experimental conﬁgurations provided by Xiong et al.
Note that due to different code implementations, as shown in Table 1, the results from our re-implemented FNet achieve 53.10 average score, lower than 55.30 reported in the original FNet paper (Lee-Thorp et al., 2021) where the original FNet is implemented in JAX/Flax.
The Pytorch implementation also causes some difference in results for other models compared to their LRA results reported in the original LRA paper (Tay et al., 2021), which are implemented in JAX/Flax.
For each task, the input sequence is truncated evenly into K segments and we experiment with
K ∈ { 32 , 64 , 128 } (as the K in Equation 4).
We ﬁnd that 32 segments for the Image task,segments for the ListOps and Retrieval tasks, 128 segments for the Text task produce the best results.
However, PoNet fails on the Pathﬁnder task under all three segment conﬁgurations.
Hence we remain computing SMP on the whole sequence (i.e., only 1 segment) for the Pathﬁnder task.
Speed and Memory Comparison
The training speed and peak memory consumption comparisons are conducted on the LRA text classiﬁcation task on a single NVIDIA Tesla V100 GPU.
The input sequence lengths are set from 512 to 16384.
Note that since the original LRA github 10 also provides the full length datasets, hence we could truncate the sequences in the LRA text classiﬁcation task into 8K and 16K lengths.
The hyper-parameters are the same as in (Xiong et al., 2021), that is, the hidden size is set to 64, the intermediate size is set to 128, the number of attention heads is set to 2, the number of layers is set to 2, and the batch size is set to 32.
A. P RE - TRAINING D ETAILS
For feasible experiment turn-around with our computation resources, we use the English Wikitext103 (100M words) 11 and the BooksCorpus (800M words) 12 datasets for pre-training.
Natural paragraphs which are marked by “ \ n” are treated as segments for the SMP computation.
For the MLM task, the masking probability is set to 15%.
80% of the masked positions are replaced by “[MASK]”, 10% are replaced by randomly sampled words, and the remaining 10% are unchanged.
For the SSO task, a long sequence containing several paragraphs is truncated into two subsequences at random positions, with 1/3 probability of replacing one of the subsequences with another randomly selected subsequence, 1/3 probability of swapping the two subsequences, and 1/3 probability unchanged.
These three cases are assigned three different labels for the ternary classiﬁcation.
All input sequences are truncated to a maximum sequence length of 512, and to accommodate sentences of different lengths, some input sequences are truncated shorter with a probability of 0.1.
The datasets were duped 5 times to alleviate overﬁtting of SSO tasks, which were also applied by Devlin et al.
Since the selection of masked positions and sentence pairs are done according to probabilities, we obtain 5 times more training sample pairs.
The pre-training implementation is based on the Pytorch codebase from Wolf et al.
(2020), with the hyper-parameters shown in Table 6.
Each pre-training experiment is run on 4 NVIDIA Tesla V100 GPUs and takes about 9 days.
[Table]
A. F INE - TUNING ON THE GLUE B ENCHMARK
The GLUE datasets 14 can be accessed from Lhoest et al.
The special token “[SEP]” is used as the segment separator.
Inputs to single-sentence tasks are processed as “[CLS] S [SEP]”; whereas inputs to sentence-pair tasks as “[CLS] S1 [SEP] S2 [SEP]”.
We implement all ﬁnetuning code based on the Pytorch codebase from Wolf et al.
We run 20 sets of hyperparameter conﬁgurations based on Table 6 for ﬁne-tuning BERT, FNet, and our PoNet and report the best GLUE results in Table 3.
Note that Table 3 provides a fair comparison between BERT,
FNet, and PoNet as all three models are pre-trained with the same pre-training data of English
Wikitext-103 (100M words) and BooksCorpus (800M words) datasets 15 , the same pre-training tasks of MLM+SSO, the same pre-training conﬁgurations (Appendix A.2)
(all trained with 340K steps), and ﬁne-tuned on GLUE with the same conﬁgurations.
A. F INE - TUNING ON THE L ONG - TEXT C LASSIFICATION T ASKS
The HND dataset can be acquired followed the guide in (Beltagy et al., 2020) 16 .
The IMDb datasetand Yelp-5 dataset 18 are from Lhoest et al.
The Arxiv-11 dataset is from He et al.
The max sequence length for all long-text classiﬁcation experiments is 4096.
Since there is no natural paragraph segmentation for these data, we use the NLTK toolkit (Bird et al., 2009) to segment [Table] the input into sentences for SMP computation.
Note that our PoNet model was pre-trained with max sequence length 512, to be able to ﬁne-tune on 4096 input lengths, following Beltagy et al.
, we add extra position embeddings initialized by copying the pre-trained 512 position embeddings recurrently.
We implement all the ﬁne-tuning code based on the Pytorch codebase from Wolf et al.
(2020) with the hyper-parameters shown in Table 6.
For BERT-Base, we ﬁne-tune using the HuggingFace BERT-Base-uncased checkpoints 20 .
We ﬁnetune FNet-Base (Lee-Thorp et al., 2021) by converting the ofﬁcial FNet checkpoints using the tool from https://github.com/erksch/fnet-pytorch to be loadable by the Pytorch codebase for ﬁnetuning.
For PoNet-Base, we ﬁne-tune our pre-trained PoNet-Base-uncased with pre-training hyperparameters shown in Table 6, that is, PoNet-Base pre-trained with 750K steps.
B A DDITIONAL GLUE R ESULTS
We also show additional GLUE validation results to compare the performance between PoNet and the ofﬁcial checkpoints of BERT and FNet in Table 7.
The results with (1) in Table 7 are from (Lee-Thorp et al., 2021).
The results with (2) in Table 7 show the best GLUE results from searching 20 sets of hyper-parameter conﬁgurations based on Table 6 for ﬁne-tuning the pre-trained models.
For BERT-Base(2) and FNet-Base(2), we use their respective pre-trained checkpoints provided by authors.
For BERT-Base(2), we ﬁne-tune on GLUE using the Huggingface BERTBase-uncased checkpoints 21 .
For FNet-Base(2), we re-evaluate FNet-Base-cased (Lee-Thorp et al., 2021) on GLUE by converting the ofﬁcial FNet-Base-cased checkpoints 22 using the tool from https://github.com/erksch/fnet-pytorch to be loadable by the PyTorch codebase for ﬁne-tuning.
For PoNet-Base(2), we ﬁne-tune our pre-trained PoNet-Base-uncased with pre-training hyperparameters shown in Table 6 and pre-trained with 750K steps.
It is important to point out that the GLUE results from ofﬁcial checkpoints of BERT-Base and FNetBase and our pre-trained PoNet-Base are not comparable since the pre-training data and pre-training tasks are different.
BERT-Base was pre-trained on English Wikipedia (2.5B words)
+BooksCorpus (800M words) with MLM+NSP, in total 20GB data size.
FNet-Base was pre-trained on C4 ( 700GB data size) with MLM+NSP.
In contrast, our PoNet-Base was pre-trained on Wikitext-103 (100M words)+BooksCorpus (800M words) with MLM+SSO, in total only 5GB data size.
Nevertheless, with the much smaller pre-training data, PoNet-Base still achieves GLUE AVG 77.54, reaching 92.84% of the accuracy of BERT on GLUE (83.52) and outperforms FNet (74.23) by 4.5% relatively, and also better than 76.7 reported in the original FNet paper (they used different code implemented with JAX/Flax). [Figure]
We also compare other reported results to verify the correctness of our ﬁne-tuning implementations.
Note that as reported in Table 7, ﬁne-tuning the FNet-Base-cased ofﬁcial checkpoints in our
PyTorch-implemented ﬁne-tuning code achieved 74.23 AVG score on the GLUE Validation set.
HuggingFace reproduced FNet-Base-cased with PyTorch and after pre-training FNet-base using the same pre-training data C4 and MLM and NSP tasks as the original FNet paper (Lee-Thorp et al., 2021) and conducting GLUE ﬁne-tuning, their FNet-Base-cased achieved 74.89 AVG score on the
GLUE validation set 23 , which is comparable to 74.23 from our FNet-Base-cased ﬁne-tuning results.
These results veriﬁed the correctness of our FNet ﬁne-tuning implementation.
HuggingFace also compares the GLUE validation results between their FNet-Base re-implemented in PyTorch and the original FNet implemented in JAX/Flax, and also shows a signiﬁcant gap from the PyTorch implementations 24 .
Hence, we think the difference between our FNet-Base GLUE score and the score reported in the original FNet paper (our 74.23 versus their 76.7) is also due to different code implementation (PyTorch versus JAX/Flax) and different hyper-parameter settings.
For the ofﬁcial
BERT-Base-uncased checkpoints
, we obtain the best GLUE validation results as AVG 83.52, which is better than the 82.62 AVG score HuggingFace reported from BERT-Base-cased 25 .
These results veriﬁed the correctness of our BERT ﬁne-tuning implementations.
C V ISUALIZATION C. S TATISTICAL A NALYSIS
To study the importance of the three granularities of pooling, i.e., GA, LMP, and SMP, at different layers of PoNet, we take 4000 examples, calculate the mean of the L 2 -Norm of each type of pooling for the 4000 examples at each layer, as well as the average of the means for the three types of pooling.
The resulting curves are shown in ﬁgure 3.
Speciﬁcally, the norm is calculated as follows,
[Figure]
[Equation]
where i, h, l, d denote the dim of example, attention head, text length, and hidden unit, respectively.
We observe that
(1) All three curves corresponding to the three types of pooling are close to the average curve in the graph, indicating that all three types of pooling play a signiﬁcant role in PoNet.
This observation is consistent with the ablation results shown in Table which shows that removing GA, removing SMP, or removing LMP in PoNet all signiﬁcantly degrade the performance on downstream single-sentence and sentence-pair tasks.
(2) All three curves corresponding to GA, SMP, and LMP share a similar trend of going down, up, down, slowly rising and going down.
From Layer 1 to 8, the norm values of SMP are greater than those of LMP and GA.
Higher than Layer 9, the norm values of LMP turn out to be greater than those of SMP.
These observations indicate that the signiﬁcance of different granularities of pooling changes across layers.
The norm value of GA is relatively low compared to that of SMP and LMP.
On the other hand, as shown in the ablation results in Table 5, removing GA from PoNet (PoNet w/o GA) degrades performance on SST and STS-B signiﬁcantly, showing that the sentencepair tasks heavily rely on the global information.
C. C ASE S TUDY
To analyze how PoNet works, we loaded a pre-trained PoNet and selected an example, “The word had ﬁlled his head as though the girl had whispered directly into both ears.”, for visualization.
For
GA, we average the attention weights of all heads as the attention weights of each layer.
For SMP, we count the number of times all hidden layer dimensions were taken as the segment maximum value.
The resulting GA attention map and SMP argmax positions across layers are shown in Figure 4.
The tokens “[CLS]” and “[SEP]” are excluded in the SMP visual map since they belong to a single word segment.
To sharpen the display, the maxima in the ﬁgure are truncated to 0.2 or 0.3 accordingly.
On the bottom layer, we observe that GA ﬁrst focuses on some of the more important words, such as “whispered” at Layer 1 and then attends to the rest words at Layer 2.
This complementary attention allows the model to capture more comprehensive information about the sequences after multiple layers.
We observe that SMP is more inclined to capture information about punctuation and pronouns at the bottom level and then some keywords, e.g., “whispered”, also begin to receive attention at the higher level, especially at Layer 9-11 based on the SMP argmax positions.
D C OMPARISON BETWEEN P O N ET AND O THER M ODELS D. D IFFERENCE BETWEEN P O N ET AND F ASTFORMER
First, our PoNet consists of three components, i.e., global aggregation(GA), segment maxpooling(SMP), and local max-pooling(LMP).
The importance of SMP and LMP is demonstrated in Section 5.
Second, the GA component in PoNet is different from Fastformer (Wu et al., 2021) due to different motivations and implementations.
In the motivation aspect, our GA is inspired by the global attention in Longformer and BigBird, but is also different from them (Appendix D.3).
In contrast, the additive attention mechanism in
Fastformer bears some similarities with the works on matrix sketching.
In the aspect of implementation, ﬁrstly, comparing GA in PoNet with Fastformer, the global query is computed differently.
In Fastformer, the global query vector is a weighted sum of all query vectors; whereas, in PoNet, it is an average pooling of all query vectors.
Secondly, the global query is mixed with key and value in different ways.
In Fastformer, elementwise product is performed between the global query vector q and each key vector k i to obtain a set of vectors p i .
A set of learnable vectors are introduced for additive attention to compute weights for p i .
After softmax normalization of the weights, the sum of weighted p i is the global key (GK).
Then, element-wise product is performed between GK and the value matrix V to obtain a set of vectors
U .
After that, U passes through the linear transformation layer to obtain its hidden representations
R .
And the hidden representations R are element-wise summed with Q to obtain the ﬁnal output of
Fastformer.
In our GA, as described in Section 3.1.1, the ﬁrst stage value for GA, denoted by g in the paper, serves as a query token to interact with the input sequence through cross-attention.
There are several differences from Fastformer.
First, standard scaled dot-production is performed on Q and K matrices, and the weight of the scalar is obtained directly.
Then, softmax is applied on the scaled dotproduct results as weights for V .
The weighted sum of V is computed as the global representation, denoted by g (cid:48) in the paper, and there is no process of global key ( GK ).
Finally, as described in
Section 3.1.4, we perform element-wise product between g (cid:48) and H on to compute the ﬁnal global representation for pooling fusion, where H on is from projecting input embedding of each token through another linear transformation speciﬁc for pooling fusion.
Different from Fastformer, there is no linear transformation layer nor summation with Q .
The above comparisons show that the GA in our PoNet is quite different from Fastformer.
In addition to algorithm comparisons, we also experimented with replacing the GA module in our PoNet with
Fastformer, and observed that after the same pre-training and ﬁne-tuning conﬁgurations, replacing
GA with Fastformer has lower performance compared to our PoNet: On CoLA the performance dropped from 45.36 to 43.28; on STS-B, the performance dropped from 84.57 to 84.46.
D. D IFFERENCE BETWEEN P O N ET AND L UNA
The mechanism in Luna (Ma et al., 2021) is also somewhat similar to our GA module.
Luna introduces an additional ”global” sequence of ﬁxed length L , which is used as a query to cross-attend with the original input to obtain a global representation of length L .
Then, it is served as key and value to perform the second cross-attention with the original input to obtain the ﬁnal output.
First, instead of introducing additional tokens to preform cross-attention, we use the results of average pooling as the global representation of the ﬁrst stage of GA.
Also, in the second stage of GA, the global representation is served as query in GA rather than key and value in Luna to calculate crossattention.
While Luna is able to directly obtain an output of the same length as the sequence,
our
GA obtains an output consistent with the length of the ﬁrst stage global representation (i.e., 1), so we further integrate this output into the original input by element-wise multiplication.
D. D IFFERENCE BETWEEN P O N ET AND L ONGFORMER /B IG B IRD
Our design of the multi-granularity pooling and pooling fusion is inspired by previous works such as Longformer and BigBird.
Our design considers the importance of capturing context information from different granularities and combining the captured information.
However, our design is different from previous works.
Longformer uses a window-based self-attention to capture local context and a task-speciﬁc global attention to encode inductive bias about the task.
For example, for classiﬁcation tasks the global attention attends to the [CLS] token and for question-answering tasks, the global attention attends to all question tokens.
In contrast, GA in PoNet is task-agnostic.
BigBird combines global, windowed local, and random attentions.
BigBird chooses a subset from existing tokens as global tokens or adds additional global tokens such as [CLS].
These global tokens attend to all existing tokens.
The GA in PoNet is computed differently from BigBird (see
Section 3.1.1 and Section 3.1.4 for details).
Also, PoNet uses GA, segment max-pooling and local max-pooling.
We experimented with replacing our windowed local max-pooling with random pooling in an early PoNet structure with a small number of pre-training steps, and observed that this change weakened PoNet’s capability of capturing local context and signiﬁcantly degraded the accuracy of downstream single-sentence CoLA task (from 42.57 to 31.44) while causing a slight gain on the sentence-pair STS-B task (from 80.29 to 80.48).
Hence, we did not use random attention in the ﬁnal structure of PoNet.
D. D IFFERENCE BETWEEN P O N ET AND FN ET
To the best of our knowledge, our work is the ﬁrst to comprehensively explore the full potential of the simple pooling mechanism for token mixing and verify that built upon multi-granularity pooling and pooling fusion, our PoNet is competitive in modeling long-range dependency with a linear complexity.
The token mixing mechanism in PoNet, built upon pooling mechanisms, is completely different from Fourier transforms used by FNet.
As shown in Table 1, on the LRA benchmark, PoNet achieves an accuracy of 61.05, which is signiﬁcantly better than the accuracy from FNet (55.30 reported in the original paper with JAX/Flax implementation, 53.10 in our PyTorch implementation).
PoNet achieves signiﬁcantly better pretraining accuracy (with the same MLM+SSO tasks) than FNet, as shown in Figure 2.
PoNet achieves signiﬁcantly better transfer learning performance on GLUE than FNet, as shown in the updated
Table 3, and signiﬁcantly better transfer learning performance on long-text classiﬁcation tasks than
FNet, as shown in Table 4.
E A DDITIONAL M ODEL C ONFIGURATIONS FOR P O N ET
We explore several additional ideas to improve PoNet
. E. T REE M AX - POOLING (TMP)
For different layers in PoNet, we apply different dilation sliding window max-pooling as another way to exchange the information between two different tokens.
We compute the TMP value T ∈ R N × d .
The size (length) of the dilation windows is based on 2, that is, on the l -th layer, len dilation = 2 l − 1 .
The lowest level, i.e. Level 1, uses a length of 1 for dilation, which is a normal sliding window max-pooling.
Since each length of the dilation sliding windows can be represented by binary and the connection state of tokens can be represented as { 0 , 1 } (not link or link), the distance between any two tokens can be reached by this structure.
We can easily calculate that the longest distance that the structure could reach was 2 l +1 − 2 .
After adding TMP into pooling fusion for PoNet, we observed that the MLM validation accuracy improves in pre-training, but we did not observe performance improvements on the downstream GLUE tasks.
E. C ONTEXTUAL D ECISION M ECHANISM
Since different tokens may need different levels of information, strictly adding the three pooling features together for all tokens, as in Equation 9, cannot meet this requirement.
We consider another aggregation method, denoted Contextual Decision Mechanism, to replace the Pooling Fusion in
Section 3.
Following the idea of attention, each token conducts a cross-attention with the 3 pooling features for contextual interactions, as follows:
[Equation]
12) where
[Equation] [Equation]
[Equation]
15) and W ∗ ∈ R d × d , b ∗ ∈ R
d
are parameters to be learned.
Note that K Mn , V Mn ∈ R 3 × d .
M (cid:48) is then the ﬁnal output of the multi-granularity pooling block.
Similar to the TMP idea, when switching from the pooling fusion in Section 3 to this contextual decision mechanism for
PoNet, we observed that the MLM validation accuracy improves in pre-training, but did not observe performance improvements on the downstream GLUE tasks.