Efficient Representation Learning via Adaptive Context Pooling

A standard transformer model (Vaswani et al., 2017) is a chain of self-attention modules (self-attention plus feed-forward layers). The input of each self-attention layer is a feature matrix $\bm{X}\in\mathbb{R}^{n\times d}$ from the preceding layer. $\bm{X}$ is a sequence of $n$ tokens =$\{\bm{x}_{1},\dots,\bm{x}_{n}\}$ each of dimension $d$. The attention layer operates on all the token features in $\bm{X}$. Specifically, each token $\bm{x}_{i}$ is first transformed to the query $\bm{q}_{i}=\bm{W}^{q}\bm{x}_{i}$, key $\bm{k}_{i}=\bm{W}^{k}\bm{x}_{i}$ and value $\bm{v}_{i}=\bm{W}^{v}\bm{x}_{i}$ with learned projection matrices $\{\bm{W}^{q},\bm{W}^{k},\bm{W}^{v}\}\in\mathbb{R}^{d\times d}$. Then the attention score of one query $\bm{q}$ attending to all the keys $\{\bm{k}_{i}\}$ stored in a memory is given by:

The final output $\bm{o}_{q}$ from querying the memory with $\bm{q}$ is obtained by taking a weighted average of all the values $\{\bm{v}_{i}\}$ in memory:

In practice, multi-head self-attention is used in transformers, where multiple projections are learned to compute attention within different heads. The outputs are then concatenated and projected into refined token features.

Drawback The above self-attention mechanism assumes a fixed granularity over which to construct the query and key vectors for individual tokens. However, such a fixed granularity may be sub-optimal for modeling context with different scales. Consider neural machine translation with word-based tokens – translating numerals from one language to another requires little context, but translating ambiguous pronouns (e.g., “it”) requires long range-contextual cues from neighboring tokens. One might argue that this difficulty can be resolved by using deeper models, where self-attention in deeper layers can capture the interactions between single tokens and bake them into deep feature representations. This can progressively correct for the fixed attention granularity at lower layers, but requires more computation that may be avoidable with an adaptive strategy.

Motivation We motivate our method using Fig. 2(a). In language modelling, if we can piece words together to form phrases, we can gradually capture the phrasal patterns and useful context information. This helps to disambiguate the pronoun “it” by linking it to the phrase “The couch”. Similarly, in image understanding, pooling similar image patches can enable the model to learn semantics of a bird’s body parts. To account for the special role of context in obtaining adaptive attention granularity, we introduce an explicit way of learning context-aware token features. We do this by learning to pool neighboring features for each token (ContextPool). Self-attention between such pooled features can thus be context-aware and model high-level dependencies, without requiring multiple self-attention layers.

Therefore, our ContextPool method needs an input-adaptive pooling function. Below, we describe how to learn that with adaptive pooling weights and pooling size. Specifically, given the input token feature matrix $\bm{X}\in\mathbb{R}^{n\times d}$, we pool for each token $\bm{x}_{i}\in\bm{X}$ with learned weights $\bm{w}\in\mathbb{R}^{n\times 1}$ and a Gaussian mask $\bm{g}^{i}\in\mathbb{R}^{n\times 1}$ (acting as a soft, local pooling window), generating a contextual feature matrix $\bm{Y}\in\mathbb{R}^{n\times d}$ of the same size of $\bm{X}$ (see Fig. 2(b)).

Adaptive pooling weights differ from the uniform ones in the popular average pooling function. We found it helpful to reweight the neighboring token features $\{\bm{x}_{j}\}$ during pooling based on their contextual support to $\bm{x}_{i}$. One widely used approach of measuring such support is based on nonlocal feature similarity as in (Wang et al., 2018):

where $\theta(\bm{x}_{i})=\bm{W}^{\theta}\bm{x}_{i}$ and $\phi(\bm{x}_{j})=\bm{W}^{\phi}\bm{x}_{j}$ are embeddings with learnable projections $\{\bm{W}^{\theta},\bm{W}^{\phi}\}\in\mathbb{R}^{d\times d}$.

We dub such learned pooling weights $\bm{w}$ as nonlocal weights (NL weights). The intuition behind NL weights is that similar features in the context are likely to correspond to semantically related entities. Therefore, nonlocal similarity pooling in form of $\sum_{i=1}^{n}w_{i}\bm{x}_{i}$ can provide contextual information to increase (or decrease) the probability of a semantic region or segment. Note we only introduce NL weights as a comparing baseline.

One limitation of NL weights is that each weight $w_{j}$ in Eq. (3) only depends on a feature pair $(\bm{x}_{i},\bm{x}_{j})$, overlooking the potential contributions from other features to $\bm{x}_{i}$. Here we turn to learning $w_{j}$ by a mapping function $m(\cdot)$ conditioned on all the token features $\{\bm{x}_{i}\}$ in $\bm{X}$. In fact, we predict the pooling weights $\bm{w}=m(\bm{X})$ all at once, where $m(\cdot)$ is implemented as two convolutional layers. Hence the prediction of $\bm{w}$ is collaborative and more efficient than NL weights prediction.

Adaptive pooling size Pooling with adaptive weights, however, does not take into account the location relationships between tokens. Here we introduce a locality prior to bias pooling towards the local context around considered token. Note that learning the pooling weights alone might also be able to find local patterns in the learned weights. However, the locality prior can simplify learning by allowing factorized and independent predictions of pooling weights and scope. Our experiments support this hypothesis with favorable results. The locality prior also shares a similar high-level idea with the effective receptive field (Luo et al., 2016), which is shown to have a Gaussian distribution.

We learn a Gaussian mask for each token to implement soft, localized pooling with adaptive pooling size rather than a hand-picked one. Specifically, we learn the mapping function $m(\cdot)$ to predict both the pooling weights $\bm{w}\in\mathbb{R}^{n\times 1}$ and sizes $\bm{s}\in\mathbb{R}^{n\times 1}$ for $n$ input tokens conditioned on their features $\bm{X}$, i.e., $\{\bm{w},\bm{s}\}=m(\bm{X})$. We implement $m(\cdot)$ again by two convolutional layers, but with the channel size set to 2 now. This enables generating the vectors of $\bm{w}$ and $\bm{s}$ altogether, which are normalized by a softmax function for ease of training. Given the normalized pooling size $s_{i}\in[0,1]$, we then transform it to the standard deviation $\sigma_{i}=rn\cdot s_{i}$ of a Gaussian mask $\bm{g}^{i}\sim\mathcal{N}(i,\sigma_{i}^{2})$. Here $r$ is a scalar empirically set as 0.1.

By multiplying the learned pooling weights $\bm{w}$ with the Gaussian mask $\bm{g}^{i}$ for token $\bm{x}_{i}$, we arrive at our final ContextPool function:

where $\bm{y}_{i}\in\bm{Y}$ denotes the ContextPooled features, $f_{ave}$ denotes average pooling function, $\gamma(\cdot)$ is a broadcasting function for element-wise multiplication $\odot$. We set the normalization factor as $C(\bm{X})=\sum_{j}w_{j}\cdot g^{i}_{j}$.

As shown in Fig. 2(c), our ContextPool module can be placed after different attention blocks. After each attention block, we take its outputs as input token features and pool them to the same number of features for use in the next attention block. During training, we jointly learn the main model and ContextPool parameters. We also show the applicability of our ContextPool method to ConvNets in supplementary materials.

Neural Machine Translation For language tasks, we first experiment on the token-level Neural Machine Translation (NMT) task. We use both the WMT 2014 English-to-German (EN-DE) dataset with about 4.5 million English-German sentence pairs, and the and English-French (EN-FR) dataset with about 36 million English-French sentence pairs. A token is a byte pair or a word piece as in (Vaswani et al., 2017). We compare with different methods all using three transformer architectures as defined in (Li et al., 2019b): Small (2 layers), Base and Big (6 layers) models. For our method, we insert ContextPool after every attention layer. Following (Li et al., 2019b), we train for 250k iterations for Small and Base models, and for 600k iterations with a smaller batch size for the Big model due to the memory constraint. We use Adam optimizer with the same learning rate schedule in (Vaswani et al., 2017).

Autoregressive Language Modeling We also evaluate ContextPool on the autoregressive language modeling task at character level. Compared to the token-level task, character-level task is harder due to much longer sequences, which would hypothetically benefit more from stronger context modeling. We use enwik8 and text8 datasets, each with 100M characters and 90M/5M/5M for train/dev/test as in (Mahoney, 2009). For testing, we follow (Beltagy et al., 2020) to split the dataset into overlapping sequences of length 32k with step size 512, and then calculate the Bits Per Character (BPC) of predicting 512 characters from previous 32k.

We use the same 12-layer model architecture with Longformer (Beltagy et al., 2020). We train our models in 3 stages with increasing sequence lengths (2048, 4096, 8192) and different batch sizes (32, 32, 16). All models are trained for a total of 530k steps with linear learning rate warmup. We also use dropout rates 0.2 and weight decays 0.01.

Image Classification We benchmark different transformer models on the widely used ImageNet-1K classification dataset (Deng et al., 2009). There are 1.28M training and 50k validation images from 1k classes. The top-1 accuracy on a single crop is reported. We consider the regular training setting in (Touvron et al., 2021) where no external training data are used. The input image resolution is $224^{2}$ by default. For higher resolutions like $384^{2}$, we fine-tune the $224^{2}$ trained models. We train for 300 epochs with the AdamW optimizer, using a cosine decay learning rate scheduler and linear warm-up (20 epochs). When fine-tuning on higher resolution images, we tune for 30 epochs with a similar training recipe as in (Liu et al., 2021). We have batch size 1024, initial learning rate 0.001, weight decay 0.05, and the max norm of gradient clipping 1. Stronger data augmentation is found to benefit our ContextPool method. Therefore we use a larger degree of augmentation with the augmentation techniques in (Touvron et al., 2021) such as RandAugment (Cubuk et al., 2020), making our pooled token features more robust.

Ablation on adaptive pooling weights and size We start with ablation studies on these two core components of our ContextPool method and compare against their alternatives. For this purpose, both the NMT and image classification tasks are considered for a comprehensive comparison. For NMT, we choose the English-German (EN-DE) translation task using the Base model. While for image classification, the ViT-B/16 model (Dosovitskiy et al., 2021) (the “Base” variant with $16\times 16$ input patch size) is used.

Tables 2 and 2 summarize the results. We observe that our ContextPool method can consistently improve the baseline transformers at only marginal overhead (in memory, FLOPs and speed), due to the efficiency of adaptive pooling functions implemented by convolutions. For the learning of pooling weights, we first compare with those un-normalized weights without using softmax (middle cell). We obtained slightly worse results for both tasks using un-normalized weights, which confirms the need of normalization for effective weighting (note the pooling size predictions were always softmax normalized). One straightforward alternative to our learned weighting is the use of uniform weights, i.e., to perform average pooling. By doing so, we save the learning cost for the weights but suffer from apparent performance loss. We can also choose to learn NL weights as in Eq. (3), which is equivalent to learning extra, single-head self-attention weights in transformers and is thus much more costly than our lightweight convolutional method. Further, NL weights are found to be less competitive than ours due to the lack of feature interactions in pairwise weights computation.

Tables 2 and 2 (bottom cell) compare several baselines to replace our learned Gaussian mask that imposes a soft locality prior for pooling. When we remove the locality prior entirely, we save compute again but observe a big drop in performance for both tasks. This suggests that context pooling indeed benefits from a local “receptive field” (similar to the findings in (Luo et al., 2016)). It also suggests the difficulty of disentangling the local prior from the pooling weights by learning the latter alone in an unfactorized way. The “Fixed window” baseline is one simple remedy to this issue by associating a fixed local window to the pooling function, where the window size is hand-picked on validation data. We see immediate help from this baseline (relative to “no locality prior”). On the other hand, we find pooling at random sparse locations will slightly hurt performance. Finally, learning adaptive local windows performs close to our method with adaptive soft Gaussian masks, but the benefits of the latter still hold with consistent gains.

Ablation on the design choice of ContextPool module Recall that our default ContextPool module is implemented as a convolutional mapping function $m(\bm{X})$, which maps the input feature matrix $\bm{X}$ into arrays of pooling weights and size. Table 3 compares such a CNN-based design choice against alternatives like fully-connected MLP and self-attention layers. Here we conduct the comparing experiments on both transformers and ConvNets (detailed in supplementary materials) for a more comprehensive ablation. Note for transformers, we still benchmark on the same tasks as in Tables 2 and 2, with identical task settings and baseline models.

The MLP-based ContextPool module in Table 3 can be considered as the simplest form of $m(\cdot)$, which maps each feature vector $\bm{x}_{i}\in\bm{X}$ to its corresponding pooling weights. We can see that MLP is more compute-efficient than our convolutional module but worse in performance. The reason is that such MLP module operates individually for $\bm{x}_{i}$ without considering feature interactions when predicting their pooling weights, while convolutional layers leverage neighboring features to do so. Note we can use a giant MLP that predicts for all $\{\bm{x}_{i}\}$ together, which becomes collaborative but at a much higher cost.

Alternatively, we can implement $m(\cdot)$ using a (single-head) self-attention layer as in Eq. (3). However, as mentioned before, such mapping function $m(\cdot)$ is not only costly with quadratic computation, but also limited in modeling feature interactions. As shown in Table 3, the inferiority also translates to the ConvNet framework. Note we can improve by modeling richer feature interactions with more than one attention layers, but this will further increase the cost.

Now we visualize what have been learned in our pooling weights and pooling sizes (in the form of soft Gaussian mask). Since visualization is easier on images with spatial grids, we take the ViT-B/16 model and visualize the predictions from our ContextPool module after the second attention layer.

We are able to observe from Fig. 3 that: 1) The pooling weights are learned to aggregate diverse information, and seem to go beyond feature similarity (the main intuition of NL weights). The last image gives one example where the pooling weights highlight some dissimilar regions around the window and ceiling, which can instead accumulate evidence for the target class of “room”. 2) The learned pooling size is indeed input dependent, capturing the local or global context adaptively. Fig. 4 further shows the distributions of pooling size in different layers. Interestingly, the predicted pooling size remains diverse within each layer, but in general tends to increase at higher layers to capture long-range dependencies.

Table 4 evaluates our ContextPool-based attention model on the token-level NMT task using both EN-DE and EN-FR datasets. Comparison is made against standard attention and other variants that model context differently. Three transformer architectures are adopted as in (Li et al., 2019b).

It is observed that local attention only achieves marginal gains over standard attention, mainly because the locality is added to the attention mechanism which hurts the full attention capacity. Area attention preserves full attention by allowing queries to attend to the whole memory. The memory is a multi-scale one to encode context of varying scales. Despite the strong BLEU scores from area attention, it is not flexible enough to model content-dependent context due to the use of fixed set of pooling sizes when constructing the multi-scale memory. Our ContextPool is able to meaningfully outperform area attention across datasets and model sizes, thanks to its adaptiveness during context pooling.

Fig. 5 further compares the above methods in terms of computation and memory complexities. Given the default number of layers $L=6$, our CP-attention not only outperforms others at the same $L$, but also strikes a better trade-off between performance and cost. For instance, our CP-attention ($L=6$) obtains a higher BLEU score 28.91 at a noticeably faster speed and lower memory than area attention, since the latter needs to maintain a multi-scale memory online. More importantly, we are able to utilize our saved compute in the form of additional layers (increasing $L$ to 7). This way, we further improve the model capacity and BLEU score, but our speed and memory remain comparable to those of area attention with only $L=6$ layers. When we continue to train a deeper model with CP, we found significantly boosted parameter efficiency over the one without CP. Interestingly, our CP-attention with $L=8$ layers obtains a similar BLEU score with the 10-layer vanilla attention model (without CP), leading to $27\%$ faster speed and $16\%$ less memory used. On the other hand, when we train shallower models, the performance gap (with vs. without CP) becomes larger, e.g., $\Delta$BLEU$=$1.21 when $L=$4. This again demonstrates our improved model expressiveness.

Finally, we evaluate on the challenging task of character-level autoregressive language modeling (see Table 5). BPC results are reported on the Dev/Test sets of enwik8 and text8 datasets. We compare with the baseline models of T12 (Al-Rfou et al., 2019) and Transformer-XL (Dai et al., 2019), as well as four representative methods of sparse attention. Among them, Adaptive local (Sukhbaatar et al., 2019) and BP-Transformer (Ye et al., 2019) use the local window as a sparsity prior, but with a learned window size and multi-scale windows respectively. Longformer (Beltagy et al., 2020) uses a combined sparsity pattern (global+local window), while Reformer (Kitaev et al., 2020) chooses to learn the patterns.

The above sparse attention methods differ from our ContextPool method in their loss of full attention capacity despite the improved efficiency. Our method on the other hand, computes full attention over ContextPooled token features. Note our context pooling function does have a locality prior, similar to existing sparsity priors based on local window. But the critical difference is that our locality prior is only applied to feature pooling, not to the following full attention process.

Table 5 confirms the benefits of full attention models. Our CP module when applied to the standard 12-layer transformer, makes a strong baseline CP-Transformer (12) that has a small model size (39M parameters) but consistently outperforms the compared sparse attention methods. We are able to further lower the model size to 36M when inserting CP to a 11-layer model without sacrificing the performance much, due to the boosted model expressiveness. When the saved parameters are re-invested in constructing a deeper model (14 layers) that has comparable model size of 44M, we attain new state-of-the-art performance on both enwik8 and text8.

The bottom cell of Table 5 examines if our ContextPool is complementary to sparse attention. The answer is positive given our consistent gains over two sparse attention baselines. Intuitively, sparse attention would benefit more from our expressive token features that are context-aware.

Table 6 evaluates our CP method on ImageNet classification and compares with the state-of-the-art methods ViT (Dosovitskiy et al., 2021), DeiT (Touvron et al., 2021) and Swin-T (Liu et al., 2021). It is shown that when CP is simply applied to the 12-layer ViT-B/16 model, performance gains are achieved at low overhead. When we plug CP into a smaller CP-ViT-B/16 model with 10 layers, this model can perform even comparably to ViT-L/16 despite being much more efficient. We further show CP is applicable to the Swin transformer that computes multi-scale attention with an image pyramid. Our CP method proves helpful for the two Swin-B models using different input image resolutions, and achieves a strong top1 accuracy of 85.6%.