PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Vision Transformers with Hierarchical Pooling

Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, Jianfei Cai

arXiv:2103.10619v2cs.CV

TL;DR

ViT models retain redundant full-length patch sequences and lack hierarchical representations, motivating a more efficient visual architecture. HVT progressively pools visual tokens to shorten sequences, supports scaling across model dimensions, and uses average-pooled tokens for prediction. With comparable FLOPs, it outperforms competitive baselines on ImageNet and CIFAR-100, although its optimal scaling strategy remains open.

  • Problem

    ViT models maintain full-length patch sequences during inference, causing redundancy and lacking hierarchical representations.

  • Method

    HVT partitions ViT blocks into stages, progressively pools visual tokens to reduce sequence length, and predicts from average-pooled visual tokens instead of relying only on a class token.

  • Results

    With comparable FLOPs, HVT outperforms the competitive DeiT baseline on ImageNet and CIFAR-100 image-classification benchmarks.

  • Takeaways & Limitations

    Hierarchical pooling improves HVT scalability, allowing depth, width, resolution, and patch size to be scaled with comparable or fewer FLOPs.

  • Takeaways & Limitations

    A principled strategy for scaling HVT to the optimal efficiency-versus-accuracy trade-off remains an open question.

Abstract

from arXiv · show

The recently proposed Visual image Transformers (ViT) with pure attention have achieved promising performance on image recognition tasks, such as image classification. However, the routine of the current ViT model is to maintain a full-length patch sequence during inference, which is redundant and lacks hierarchical representation. To this end, we propose a Hierarchical Visual Transformer (HVT) which progressively pools visual tokens to shrink the sequence length and hence reduces the computational cost, analogous to the feature maps downsampling in Convolutional Neural Networks (CNNs). It brings a great benefit that we can increase the model capacity by scaling dimensions of depth/width/resolution/patch size without introducing extra computational complexity due to the reduced sequence length. Moreover, we empirically find that the average pooled visual tokens contain more discriminative information than the single class token. To demonstrate the improved scalability of our HVT, we conduct extensive experiments on the image classification task. With comparable FLOPs, our HVT outperforms the competitive baselines on ImageNet and CIFAR-100 datasets. Code is available at https://github.com/MonashAI/HVT

1. Introduction

Transformers offer strong representation through self-attention but remain computationally expensive, while ViT models also retain redundant full-length patch sequences and lack hierarchical representations. HVT addresses these issues through hierarchical pooling and visual-token averaging, improving scalability and classification performance with comparable FLOPs.

  • Motivation: Self-attention enables strong long-range dependency modeling, but Transformer models incur high computational costs that hinder deployment on resource-constrained devices.The cited motivation notes that a single Transformer model can require more than 10G Mult-Adds for translating a 30-word sentence.
  • Motivation: ViT models maintain full-length patch sequences across layers, creating redundancy and lacking the multi-level hierarchy important for image recognition.The authors identify layer-dependent redundancy and the absence of hierarchical representations as two limitations of conventional ViT designs.
  • Proposed approach: HVT partitions Transformer blocks into stages and progressively pools visual tokens to shrink sequence length, paralleling hierarchical spatial downsampling in CNNs.Pooling is applied as the model goes deeper, reducing sequential resolution while constructing a pyramidal representation.
  • Proposed approach: HVT replaces class-token-only prediction with average pooling over visual tokens, which empirically provides richer discriminative patterns for classification.The authors report that average-pooled visual tokens contain more discriminative information than the single class token.
  • Results: With comparable FLOPs, HVT outperforms DeiT on ImageNet and CIFAR-100 image-classification benchmarks.The contribution summary presents this as the main empirical result across the two datasets.

2. Related Work

Prior work improves Transformer efficiency through compression, architectural search, or reduced-complexity attention, while visual Transformer research spans hybrid convolution-attention and pure-attention designs. HVT instead applies hierarchical pooling to ViT sequences to reduce computation and improve scalability across key architectural dimensions.

  • Visual Transformers: Visual Transformer research includes hybrid architectures that combine convolutional backbones or layers with self-attention for computer-vision tasks.Examples include DETR, BotNet, and other convolution-attention combinations.
  • Visual Transformers: Another line of work designs pure attention-based visual architectures by replacing spatial convolutions or local aggregation mechanisms with self-attention.The related-work discussion includes models that replace convolutional operations with global or local attention mechanisms.
  • Efficient Transformers: Efficient Transformer methods reduce resource use through quantization, pruning, distillation, Neural Architecture Search, kernelization, sparsity, or low-rank attention.These approaches target either generic model compression or the quadratic complexity of self-attention.
  • HVT positioning: HVT targets coarse-patch ViT limitations with hierarchical pooling that reduces computational cost while scaling depth, width, resolution, and patch size.The strategy also aims to preserve a pyramidal feature hierarchy for visual representations.

3. Proposed Method

HVT adapts ViT into a hierarchical encoder by progressively pooling visual tokens, updating positional information after sequence reduction, and predicting from average-pooled patch tokens rather than a class token. This design reduces computational cost while retaining a sequence-based representation.

  • Overview: The proposed HVT framework is intended to reduce full-length sequence redundancy and construct a hierarchical visual representation.The paper presents hierarchical pooling and class-token-free prediction as its two main components.
  • Hierarchical Pooling: HVT partitions Transformer blocks into stages and applies 1D pooling to progressively shrink the visual-token sequence.The design is inspired by hierarchical VGG- and ResNet-style networks.
  • Hierarchical Pooling: After pooling changes sequence length, HVT adds stage-specific learnable positional embeddings before forwarding tokens to the next block.The additional embeddings compensate for the original positional encoding becoming unsuitable after pooling.
  • Prediction without the Class Token: HVT removes the class token and predicts from the average-pooled final patch-token sequence through normalization and a fully connected layer.The method uses the remaining sequence rather than relying on a single class token.
  • Complexity Analysis: Pooling reduces Transformer block cost when sequence length is halved, with the compression ratio bounded by α ∈(2, 4).The ratio is derived from the FLOPs of blocks operating on sequence lengths n and n/2.

4. Discussions

The discussions connect hierarchical pooling with CNN-like feature hierarchies and show that reduced computational cost supports scaling HVT across model dimensions. However, a principled method for selecting configurations remains open.

  • Feature Hierarchy: Hierarchical pooling gives HVT a pyramidal feature hierarchy analogous to the progressive downsampling used in CNNs.The discussion compares visualized feature maps from ResNet, DeiT-S, and HVT-S-1.
  • Scalability: Four pooling operations reduce DeiT-S computation by nearly 3.3× from its approximately 4.6G FLOPs baseline.The example uses a 224×224 image, patch size 16, 12 blocks, and 6 heads.
  • Scalability: The saved computation can be reallocated to wider or deeper HVT variants while keeping computational cost below DeiT-S and increasing model capacity.The discussion gives wider models with 11 heads and deeper models with 48 blocks as examples.
  • Open Question: Finding a principled way to scale HVT for the optimal efficiency-versus-accuracy trade-off remains an open question.The current exploration evenly partitions blocks and follows DeiT settings for fair comparison.

5. Experiments

Experiments evaluate HVT against DeiT and PoWER on ImageNet and CIFAR-100, examining efficiency, scaling, pooling, token prediction, and architectural choices. HVT reduces computation through hierarchical pooling, while scaling and design choices determine when accuracy improves or degrades.

  • Main Results: Experiments compare HVT with DeiT and PoWER on ImageNet, using Top-1/Top-5 accuracy, FLOPs, and parameter count.The evaluation includes tiny and small HVT settings, architectures with different pooling stages, and scaled variants.
  • Main Results: Nearly 2× FLOPs reduction from hierarchical pooling causes accuracy degradation, including a 1.80% Top-1 drop for HVT-S-1 and a 2.56% drop for HVT-Ti-1.The smaller drop for HVT-S-1 is attributed to its greater parameter redundancy.
  • Main Results: HVT-Ti-1 uses less FLOPs than PoWER while achieving better performance, whereas HVT-S-1 reduces more FLOPs but performs slightly worse.PoWER uses three training steps, while HVT uses a simpler one-stage training scheme.
  • Main Results: 3.03% higher ImageNet Top-1 accuracy is achieved in the tiny setting when saved FLOPs are reallocated to scale HVT under similar computational cost.Scaling enlarges embedding dimensions and adds attention heads, producing a better solution according to the reported experiments.
  • Ablation Study: Average pooling over patch tokens outperforms prediction with a class token, with larger gains on CIFAR-100 than ImageNet.The paper attributes the improvement to discriminative information retained across the entire token sequence.
  • Ablation Study: Pooling stages improve CIFAR-100 performance while reducing FLOPs, but increasing stages degrades ImageNet accuracy.The authors suggest pooling acts as a regularizer on redundant CIFAR-100 representations, whereas ImageNet may require better scaling.

6. Conclusion and Future Work

The paper presents HVT for image classification, using hierarchical pooling to reduce computational cost and improve scalability. It also finds that visual tokens are more informative for prediction than a single class token, while identifying encoder-only scope and future extensions.

  • HVT uses hierarchical pooling to compress sequential resolution and reduce computational cost.The method progressively pools visual tokens as depth increases.
  • The reduced computation enables scaling depth, width, resolution, and patch size with comparable or fewer FLOPs.Saved computation can be reallocated to increase model capacity.
  • Visual tokens contain more important information for class prediction than a single class token.The paper reports this result empirically for HVT image classification.
  • The paper targets HVT as an encoder for image classification, leaving decoder extensions and other vision tasks for future work.Suggested tasks include object detection and semantic or instance segmentation.

Appendix

The appendix organizes supplementary material into explanations of Transformer components and FLOPs computation details.

  • Section S1 explains the components of a Transformer block, including multi-head self-attention and the position-wise multilayer perceptron.
  • Section S2 provides details for calculating the FLOPs of a Transformer block.

S1.1. Multi-head Self-Attention

The multi-head self-attention description derives attention from query, key, and value projections, then combines head outputs through projection.

  • Self-attention computes query, key, and value matrices from the input using linear transformations.The input has sequence length N and embedding dimension D.
  • The attention map is formed from scaled inner products of queries and keys and normalized with softmax.The attention map has shape N×N and represents pairwise token scores.
  • The attention operation applies the attention map to value vectors to produce an output matrix.
  • Multi-head attention concatenates the individual head outputs and applies a linear projection.The projection uses a learnable parameter matrix.

S1.2. Position-wise Multi-Layer Perceptron

The position-wise multilayer perceptron is described as part of the Transformer block and uses two fully connected layers with GELU nonlinearity.

  • The MLP contains two fully connected layers with a GELU nonlinearity.

S2. FLOPs of a Transformer Block

The section derives FLOPs for multi-head self-attention and MLP layers, then combines them into the total cost of one Transformer block as a function of sequence length n and embedding dimension d.

  • MSA layer: An MSA layer costs 4nd^2 + 2n^2d FLOPs after combining QKV projections, attention-map computation, self-attention, and output projection.The four components are 3nd^2, n^2d, n^2d, and nd^2, respectively.
  • MLP layer: The MLP layer uses two fully connected projections, mapping each token from R^d to R^4d and back, for a total of 8nd^2 FLOPs.Each fully connected layer contributes 4nd^2 FLOPs.
  • Transformer block: Combining the MSA and MLP costs gives 12nd^2 + 2n^2d FLOPs for one Transformer block.The block cost is the sum of the two layer-level expressions.
Loading 2103.10619v2…