PnP-DETR: Towards Efficient Visual Analysis with Transformers

Introduction

Object detection is a fundamental computer vision task aiming to recognize object instances in the image and localize them with precise bounding boxes. Modern detectors address this set prediction task mainly with proxy learning objectives, i.e., regressing offset from pre-defined anchor boxes or boundaries from grid locations . Those heuristic designs not only complicate the model design but also require hand-crafted post-processing for duplicate removal. A recent method DETR eliminates those hand-crafted designs and achieves end-to-end object detection. It builds an effective set prediction framework on top of convolution feature maps with transformers and shows competitive performance to the two-stage Faster R-CNN detector. The image feature map is flattened in the spatial dimension into one-dimensional feature vectors. The transformer then processes them with its strong attention mechanism to generate the final detection list.

Albeit simple and effective, applying the transformer networks to a image feature map can be computationally costly, mainly due to the attention operation over the long flattened feature vectors. These features may be redundant: natural images often contain enormous background areas apart from the interested objects, which may occupy large part in the corresponding feature representation; also, some discriminative feature vectors may already suffice for detecting the objects. Existing works improving the transformer efficiency mainly focus on accelerating the attention operation , and few of them consider the spatial redundancy discussed above.

To address the above limitation, we develop a learnable poll and pool (PnP) sampling module. It aims to compress an image feature map into an abstracted feature set composed of fine feature vectors and a small number of coarse feature vectors. The fine feature vectors are deterministically sampled from the input feature map to capture the fine foreground information, which thus are crucial for detecting the objects. The coarse feature vectors aggregate information from the background locations, and the resulting contextual information helps better recognize and localize the objects. A transformer then models the information interaction within the fine-coarse feature space and obtains the final result. As the abstracted set is much shorter than the directly flattened image feature map, the transformer computation is reduced significantly and mainly distributed over the foreground locations. Our approach is orthogonal to the approaches improving the transformer efficiency and can be further combined with them to obtain more efficient models.

Concretely, the PnP module is composed of two core sub-modules: a poll sampler and a subsequent pool sampler. The poll sampler incorporates a content-aware meta-scoring network that learns to predict the infromativeness score of the feature vector at each spatial location. The feature vectors are then ranked spatially with the informativeness scores and a subset of most informative feature vectors are selected. The subsequent pool sampler dynamically predicts attention weights on the non-sampled feature vectors and aggregates them into a small number of feature vectors that summarize the background information. Similar to the region proposal networks , the PnP module also aims to extract object-relevant information, but is end-to-end learned without explicit objective like object bounding box regression. We build a PnP-DETR with the PnP module, which operates on the fine-coarse feature space and adaptively allocates its transformer computation in the spatial domain. Fig. 1 is an example detection with computation density map (refer to Sec. 4.2 for details of the map construction). Existing methods of improving model efficiency still need train multiple models of different complexities for achieving various trade-offs of computation and performance. Compared with them, the proposed PnP sampling allows the transformer to work with a variable number of input feature vectors and achieve instant computation and performance trade-off.

We conduct extensive experiments on the COCO benchmark, and the results show PnP-DETR effectively reduces the cost and achieves dynamic computation and performance trade-off. For example, without bells and whistels, a single PnP-DETR-DC5 obtains a 42.7 AP with 72% reduction of transformer computation compared to the 43.3 AP baseline and competitive 43.1 AP with 56% reduction. We further validate the efficiency gain with panoptic segmentation and the recent vision transformer model (ViT ). For example, PnP-ViT achieves near half of FLOPs reduction with only 0.3 drop of accuracy. To summarize, the contributions are:

We identify the spatial redundancy issue of the image feature map, which causes excessive computation of the transformer network in a DETR model. We therefore propose to abstract the feature map, so as to significantly reduce the model computation.

To realize the feature abstraction, we design a novel two-step poll-and-pool sampling module. It first employs a poll sampler to extract the foreground fine feature vectors, and then utilizes a pool sampler to obtain the contextual coarse feature vectors.

We then build PnP-DETR, wherein the transformer operates on the abstract fine-coarse feature space and adaptively distributes the computation in the spatial domain. PnP-DETR is more efficient and achieves instant computation and performance trade-off with a single model, by varying length of the fine feature set.

The PnP sampling module is general and end-to-end learned without explicit supervision like the region proposal networks . We further validate it on panoptic segmentation and recent ViT model and show consistent efficiency gain. We believe our method provides useful insights for future research on efficient solutions of vision tasks with transformers.

Related Work

In recent years performance of object detection has been substantially improved over traditional approaches . Those modern methods mainly address the task with a relaxed learning objective, i.e., learning on a set of matched positive anchor box samples and predicting with post-processing (NMS) to suppress duplicates. The handcraft designs Recently, proposed an end-to-end DETR framework that learns an explicit set based objective with transformers , showing decent performance compared to previous two-stage methods . Our work aims at improving efficiency of end-to-end objectors by reducing spatial redundancy. Compared to most recent deformable DETR that improves the attention efficiency, we aim to directly compress the feature map, which is from different perspective and could be potentially combined together. For example, by implementing bilinear interpolation kernel in the irregular sampled space to enable the learning of deformable offset prediction.

Sparse Execution and Sampling

Lots of works explored sparse execution in convolution layers , saving computation by avoiding convolution operations on some less informative spatial locations. In this work, we are partially inspired by the sparse convolution and explore sparse execution of transformers by developing a dynamic image feature sampling method for efficient subsequent processing. Our work is also related to literature on learning a sampling policy for point cloud understanding tasks . Different from these works where sampling is achieved by new data point generation, we directly address discrete sampling by using a novel sampling as ranking strategy.

Method

Without loss of generality, DETR first utilizes a backbone convolution network $\mathcal{C}$ with parameters $\theta_{c}$ to extract the image feature map F:

$(cls_{k},box_{k})$ denotes one detected object with category and bounding box, the number of detections is fixed to $D$.

2 Feature Abstraction

We propose a feature abstraction scheme to address the above limitation. It obtains two sets of feature vectors for compact feature representation:

3 Poll and Pool (PnP) Sampling

The larger the score is, the more informative the feature vector $\textbf{f}_{ij}$ is. We then sort all the scores $\{s_{ij}\}$ as

where $\aleph$ is the sorting order and $L=HW$. With $\aleph$, we then take the top $N$ scoring vectors to form the fine feature set:

We use layer normalization and turn off the affine parameters. Ideally, $N$ may vary with the image content, but we observe that fixed amount sampling already generates good performance, i.e., $N=\alpha L$ where $\alpha$ is a constant fractional value, which we name as the poll ratio. This design also enables an extension to single model computation and performance trade-off discussed in Sec. 3.4.

Pool Sampler

The above poll sampler extracts the fine feature set. The remaining feature vectors mainly correspond to the background area. To compress them into a small feature set that summarizes the contextual information, we design a pool sampler that performs a weighted pooling of the remaining feature vectors to obtain a fixed number of $M$ background contextural feature vectors. This is partially inspired by the bilinear pooling and double attention operation where global descriptors are generated for capturing the second-order statistics of the feature map. Formally, the remaining feature vector set is

We then normalize the aggregation weight over all the remaining non-sampled locations with softmax:

With the normalized aggregation weight, the projected feature vectors are aggregated to obtain a new feature vector that summarizes the information of non-sampled locations:

By aggregating with all $M$ aggregation weights, we obtain the summarized coarse background contextual feature set:

Reverse Projection for Dense Prediction Tasks

The PnP module reduces the image feature map from 2D coordinate space to an abstracted space, which cannot be used for dense prediction tasks like image segmentation. To address the limitation, we propose to project the encoder output feature vectors back to the 2D coordinate space. Specifically, the fine feature vectors are scattered back to the sampled locations; the coarse feature vectors are first diffused back to original 2D space with the aggregation weight:

and then scattered back the non-sampled locations of the poll sampler. $\hat{\textbf{f}_{m}}$ denotes output coarse feature vector from the encoder and $\hat{\textbf{f}}_{r}$ means the projected feature vector. The obtained 2D feature map is then used for dense prediction.

4 PnP-augmented Models

The PnP module is general and straightforward. It can be plugged into existing models to enable them to operate on the fine-coarse feature space for better efficiency. We here describe the models we build to evaluate the PnP module and our proposed random poll ratio scheme to enable instant computation and performance trade-off with a single model.

Recently introduced a transformer-based image recognition model named Vision Transformer (ViT). We evaluate the generalizability of our method on the ViT model. We build the PnP-DETR and PnP-ViT by plugging the PnP module before the transformer network. The resulting models are end-to-end learned and other settings are the same with original models. We use the hybrid ViT architecture . Unlike original DETR and ViT wherein the transformer directly operates over the full image feature space, the PnP augmented transformer models the information interaction on the fine-coarse feature space and adaptively allocates its computation in the spatial domain to achieve better efficiency.

Instant Computation and Performance Trade-off

To achieve different computation and performance trade-offs, existing methods improving transformer efficiency generally train multiple models with different complexities controlling hyperparameters, e.g., number of hashes in Reformer and projected feature dimension in Linformer . Unlike them, a model equipped with a PnP module can achieve instant single model computation and performance trade-off. This is enabled by controlling the poll ratio $\alpha$ to determine the amount of fine information preserved. With a larger $\alpha$, more fine feature vectors are obtained, and the overall performance is expected to be higher; with a smaller $\alpha$, the performance may be lower but more computation is saved. However, we find inference with a different $\alpha$ to training severely degrades the performance. We propose to generate a random poll ratio during training:

Where $\alpha_{low}$ and $\alpha_{high}$ defines the value range. $\alpha$ is updated in each iteration. In this way, the transformer learns to work with variable length of input feature vectors, and thus achieves the desired single model computation and performance trade-off by inferring with different poll sample ratios (Fig. 3). The model only needs to be trained once.

Experiments

For training PnP-DETR, we use 4 images per GPU on 8-GPU machine, with a total batch size of 32. For training PnP-ViT, we use 32 images per GPU, with a total batch size of 256. The meta-scoring network is instantiated with a 2-layer MLP. Unless otherwise stated, the pool sample number $M$ is set to 60 and 240 for R50 and R50-DC5 models, respectively. Other settings including hyper-parameters, network architecture and loss functions follow the baselines for fair comparison. Due to space limit, we defer more details like position embeddings to supplementary.

2 Experiments on Object Detection

Tab. 1 shows the results of the fixed poll ratio training on the COCO benchmark. For the DETR-R50 model, with an $\alpha$ = 0.33, PnP-DETR achieves 41.1 AP and 60% reduction of transformer computation cost. Further increasing $\alpha$ to 0.5, the performance reaches a similar level as the DETR baseline (AP of 41.8 vs. 42.0), with 45% reduction of the computation. For DETR-R50-DC5 model, a similar trend is observed but more computation is saved. We also evaluate the setting of mismatched training and test poll ratio. The model trained with $\alpha$ = 0.33 gets nearly 5 AP drop when evaluating with $\alpha$ = 0.5. This observation shows the necessity of applying random poll ratio training for the model to work with variable poll ratio. We also compare to the deformable DETR , as we did not incorporate multi-scale features, which is not the focus of this work, we compare to single scale deformable DETR for fair comparison. Our method performs better than deformable DETR with less FLOPs, especially for large objects, e.g., APl of 60.0 vs. 57.8 for the ResNet-50 backbone.

Dynamic Poll Ratio Training

As shown in Fig. 4, by training with the random poll ratio with a value range of $(0.15,0.8)$, the obtained model can achieve dynamic computation and performance trade-off by evaluating with variable poll ratio. The AP for certain poll ratio is similar to the fixed poll ratio trained counterpart. For example, a PnP-DETR-R50 model gets 41.1 AP with fixed poll ratio 0.33 training and 41.2 AP with random poll ratio training. The performance is the same to the baseline with a poll ratio of 0.65. We observe when the poll ratio is large, e.g., 0.5, increasing the poll ratio brings diminished gain in AP. This is likely because the fine feature set already covers the essential spatial locations for detecting the objects, and thus more fine information only brings limited gain. Similar observations are made with the ResNet-101 backbone. Tab. 2 shows the inference time compared to baseline model, the inference time is significantly reduced.

Visualization of Computation Density Map

Fig. 5 shows some example detection results and associated computation density maps, with poll ratio of 0.33. The objects are well detected while the computation is dynamically allocated to the spatial domain in a content-aware manner. To compute the density map, we assign a weight to each spatial location. For poll sampled locations, the weight is 1. For each of other locations, the weight is the cumulative value of all pool sample aggregation weights at this location. Then the transformer cost is distributed with the normalized weights to obtain the computation density map.

3 Experiments on Other Tasks

Following , we evaluate our method on the panoptic segmentation task. To perform dense per-pixel segmentation as DETR, we project the encoder output feature back to the original 2D coordinate space. As shown in Tab. 3, the model saves computation and achieves instant performance and computation trade-off by varying the poll ratio $\alpha$, e.g., achieving Panoptic Quality (PQ) of 43.2 compared to 43.4 of a baseline DETR model, with 5G less FLOPs (i.e., 6.6G vs. 11.6G).

Image Recognition

We also apply the PnP sampling to the recent transformer-based image classification model of ViT . We use the hybrid architecture with ResNet50-stage4 feature map (14x14) and train the model on the ImageNet-1k dataset from scratch. We set the pool sample number to 10. We train the PnP-ViT with random poll ratio in the value range of $[0.2,0.8]$. As shown in Tab. 4, the PnP-ViT achieves dynamic computation and performance trade-off as observed with the DETR model. The results show the generalizability of PnP sampling design.

4 Model Analysis

We then provide several experimental analysis to better understand the proposed method. To save experiment time, we sample the COCO benchmark to obtain a smaller dataset and conduct all experiments on the sampled COCO dataset. We design a class-incremental sampling that helps preserve the data distribution. Due to space limit, we defer sampling details and more experiments to supplementary.

As shown in Fig. 6, we vary the the poll sample ratio and the pool sample number to obtain the performance curve with the same amount of computation cost. We observe that 1) with only poll sampling ($\alpha$-0.4), the performance is suboptimal; incorporating pool feature vector samples can significantly improve AP with the complementary background information from non-sampled locations, e.g., $\alpha$-0.39-$M$-10 model achieving about 0.7 AP higher than the $\alpha$-0.4 model. 2) with only pool sampling, the performance drops by a large margin. We assume it is difficult for the pool sampling to preserve accurate fine information, as it is designed to aggregate feature vectors spatially from different locations. 3) the optimal setting is $1/3$ poll ratio with $60$ pool samples, indicating that a compact feature set should be mainly composed of fine feature vectors for accurate object detection. We further individually examine the effects of pool sample number $M$ and poll sample ratio $\alpha$: 1) We vary $M$ by fixing $\alpha$. 2) We vary $\alpha$ by fixing $M$. Due to space limit, we defer the experiment results and analysis to the supplementary.

Visualizing Poll and Pool Sampling

As shown in Fig. 7, we visualize the poll sampler’s scoring map, its sampled locations, and example aggregation weight map of the pool sampler. To summarize, 1) the poll sampler learns to sample the locations within and surrounding objects; 2) the pool sampler obtains different scales of context. For example, on the first row, the first pool sample attends to a wide range of spatial locations and encodes global context information; the second sample attends to a small area around the sky, and thus captures local context. We also have some other intriguing observations on the poll sampler: 1) It learns to sample object alike area beyond the object categories used for training. For example, for the last row in Fig. 7, locations around the traffic signs and the tree-like object are sampled. The behavior is similar to a learned region proposal network (RPN) , but learned without explicit supervision. 2) It tends to sample coarsely for some large and ‘easy’ objects but finely for small ones. For example, fewer points are sampled for the woman in the first row and the bed in the second row; the books in the second image and the cars in the last image are smaller and more difficult to detect, so the poll sampler finely samples feature vectors for those objects and surrounding areas.

Tracking Poll Sampler Learning

To better understand the learning process and dynamics of the poll sampler, we record two statistics during training: (1) the proportion of sampled locations that are within the GT bounding boxes; (2) the pixel IOU of the sampled locations between consecutive epochs. As shown in Fig. 8, we make following observations. 1) The poll sampler gradually learns to sample more feature vectors that lie within the ground truth area but finally remains steady at about 60%, indicating that it also attends some background and contextual locations that are crucial for recognizing and detecting the objects. 2) The poll sampler initially has a large variation on its sampled locations, and thus the sampled areas of consecutive epochs have small IOU (i.e., about 0.2). During training, the IOU quickly converges to about 0.7 with around 30 epochs and remains steady at about 0.75, indicating that the sampler quickly learns to sample crucial feature vectors and the sampled locations does not change much. After learning rate decay at 100 epoch, the IOU of the consecutive epoch is close to 1.0, meaning the poll sampler converges.

Conclusion

In this paper, we encapsulate the idea of reducing spatial redundancy into a learnable PnP module. It is composed of a ranking based poll sampler that discretely samples fine feature information and a subsequent adaptive pool sampler that summarizes the background contextual information. The PnP module is general and can be incorporated into existing model for efficient processing while maintaining the performance, which is verified on object detection, panoptic segmentation and image recognition. We believe the proposed method offers insights for future research into efficient visiual analysis with transformers.

References

Computation Saving

Here we show the concrete computation saving by the abstraction scheme, assume the length of the full feature set is $L=HW$ and the fraction of abstracted feature length is $r=(N+M)/L$. As shown in first row of Tab. 5, for encoder, since the complexity of self-attention layers is $\mathcal{O}(L^{2})$ and the complexity of other layers (projection layers, feed-forward layers, normalization, e.t.c) is $\mathcal{O}(L)$,we assume their actual computation cost is $aL^{2}$ and $bL$ correspondingly. For the decoder, since the complexity of cross-attention is $\mathcal{O}(L^{2})$, and the complexity of other parts is not related to the sequence length $L$, we assume their costs are $cL$ and a constant $O$ respectively.

With a larger sequence length $L$ the rate is more close to $r^{2}$ and more computation is saved.

The total computation of decoder compared to original is

With a larger sequence length $L$ the rate is more close to $r$ and more computation is saved.

More Implementations

Here we describe the implementation details about padding masks and position embedding. For the fine feature set, we use the same sampling order of poll sampler to gather the corresponding position embeddings and padding masks. For the coarse feature set, we set the masks to $False$ to indicate that they are not paddings and employ pseudo position embedding by linearly combining position embeddings of the remaining feature set with the aggregation weight.

In this section, we present the detailed about how we sample the COCO dataset to obtain a smaller version for faster experimental validation. The COCO dataset has a skewed distribution of training image number over object categories, i.e., some categories have significantly smaller number of training images. Direct random sampling on all training images may cause too much loss of images on those scarce categories and the overall distribution may be even more biased. The mAP result on the biased dataset may be unstable and cannot well evaluate the model performance. To curcumvent the difficulty and obtain more effective sampled dataset, we design a new strategy. We rank the object categories according to their training image number, then perform an incremental sampling starting from the most scarce category to the most abundant category. The the sampling algorithm is given in Algorithm 1. Concretely, for each category, if the number of training images is more than a sampling threshold number and the number of already sampled images for this category is less than the threshold number, then a sampling will be performed to obtain additional training images for reaching the threshold number. As shown in Fig. 9 is the distributions of obtained sampled versions of the COCO dataset, with different setting of the sampling threshold. The sampled dataset will be smaller given a smaller threshold. We use a sampling threshold of 500 to obtain a sampled COCO and conduct all the ablation experiments on the dataset. With the designed incremental sampling, the distribution of training images over most object categories is roughly uniform, and thus can be used to more stablly evaluate model performance than a randomly sampled sub-dataset while saving enormous experiment time.

Additional Ablations

To individually examine the effect of $M$ and $\alpha$, we conduct following experiments: 1) varying $M$ by fixing $\alpha$. As shown in Tab. 6, compared to the model with only poll sample feature vectors ($M$-0), adding 30 pool feature vectors gets about 1 AP improvement, but when $M$ is larger than a certain value, the improvement is diminished (i.e., $60$). This phenomenon indicates that a small number of summarized feature vectors for the background contextual information is enough. 2) varying $\alpha$ by fixing $M$. As shown in Tab. 7, when the poll ratio $\alpha$ is small, increasing it significantly improves the performance (e.g., 25.2 AP to 27.1 AP by increasing $\alpha$ from 0.1 to 0.2). This observation shows the importance of fine information for detecting the objects. When $\alpha$ is larger than about 0.5, the performance improvement is diminished, which is as expected since the feature vectors that rank lower mostly correspond to the background locations, and thus the gain from including fine information on those locations is small.

Different Architecture of Scoring Network

As shown in Tab. 8 is the result of different network architecture of the scoring network of the poll sampler, increasing the layer number from 1 to 2 improves the AP by 0.8 (i.e., 1-layer-fc and 2-layer-fc-256.). This is likely because the 2-layer network much more accurately predict the informativeness score. Further increasing the layer number gives diminished gain, i.e., 28.8 vs. 28.7 AP for 3-layer-fc-256 and 2-layer-fc-256. We also tried decreasing the hidden neuron unit number from 256 to 32, which reduces the computation, but the performance decreased, i.e., 28.2 for the 2-layer-fc-32 scoring network, which is 0.6 lower than the 2-layer-fc-256 network in AP. We choose the 2-layer-fc-256 network as the default architecture of the score network.

Pool Sampler on The Full Feature Set

While the proposed pool sampler operates on the non-sampled feature vectors of the poll sampler, it is interesting to see if directly applying the pool sampler on the full feature set for generating the coarse feature set would be better. As shown in Tab. 9, such setting leads to about 0.5 AP drop compared to the proposed two-step setting. This may be caused by the redundant information that have been captured by the fine feature vectors from polled samples.

Comparing the Proposed Sampling Strategies to Some Alternative Methods

We compare the proposed poll sampler to some baseline alternatives including 1) random sampling: for each image, randomly sample the same amount of locations as the poll sampler and fix the sampled locations for training and evaluation. 2) uniform grid sampling: uniformly sample the 2D locations with equal interval. We adopt a general sampling mapping of $\lfloor\frac{i}{\sqrt{r}}\rfloor\lfloor\frac{j}{\sqrt{r}}\rfloor,i=0,1,...,\lfloor W*\sqrt{r}\rfloor,j=0,1,...,\lfloor H*\sqrt{r}\rfloor$ ($H$,$W$ are the height and width of the feature map and $r$ is the sampling ratio). With some specific poll ratio, the sampling is equavalent to MaxPooling, e.g., $r=1/4$ is equavalent to MaxPooling with kernel size 1 and stride 2. 3) direct interpolation: use interpolation to directly resize the feature map to target size $(\lceil H/\sqrt{r}\rceil,\lceil W/\sqrt{r}\rceil)$. As shown in Tab.10, compared to proposed ranking based poll sampling, random sample leads to a large drop in AP, i.e., 22.9 vs 27.3 for the without pool sampling setting and 23.9 vs 28.7 for the with pool sampling setting. Uniform grid sampling and direct interpolation also generate lower performance than poll sampling, e.g., 25.9 and 26.1 compared to 27.3 under the without pool sample setting. The result shows the proposed poll sampler learns effective sampling policy and is better than those simple baselines.