Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding

Zizhao Zhang; Han Zhang; Long Zhao; Ting Chen; Sercan O. Arik; Tomas Pfister

Source-linked AI summary

Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding

Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan O. Arik, Tomas Pfister

arXiv:2105.12723v4May 26, 2021cs.CV

TL;DR

Vision transformers often need sophisticated hierarchies and massive datasets, motivating a simpler and more data-efficient design. NesT nests local transformers over image blocks with hierarchical aggregation, achieving strong classification and generation results while enabling GradCAT visual interpretation.

Problem
Vision transformers can require hundreds of millions of images for pre-training, making their data requirements impractical in some settings.
Method
NesT nests canonical transformers operating on non-overlapping image blocks and connects them hierarchically with a block aggregation function.
Results
NesT improves data efficiency and convergence across ImageNet and small datasets, extends to a strong image-generation decoder, and supports GradCAT interpretability.
Takeaways & Limitations
A simplified nested hierarchy can support efficient visual representation learning, image generation, and tree-based visual interpretation.

Abstract

from arXiv · show

Hierarchical structures are popular in recent vision transformers, however, they require sophisticated designs and massive datasets to work well. In this paper, we explore the idea of nesting basic local transformers on non-overlapping image blocks and aggregating them in a hierarchical way. We find that the block aggregation function plays a critical role in enabling cross-block non-local information communication. This observation leads us to design a simplified architecture that requires minor code changes upon the original vision transformer. The benefits of the proposed judiciously-selected design are threefold: (1) NesT converges faster and requires much less training data to achieve good generalization on both ImageNet and small datasets like CIFAR; (2) when extending our key ideas to image generation, NesT leads to a strong decoder that is 8$\times$ faster than previous transformer-based generators; and (3) we show that decoupling the feature learning and abstraction processes via this nested hierarchy in our design enables constructing a novel method (named GradCAT) for visually interpreting the learned model. Source code is available https://github.com/google-research/nested-transformer.

Introduction

NesT addresses vision transformers’ data inefficiency and complex local-attention designs by nesting basic block transformers with hierarchical aggregation. The resulting architecture improves data efficiency, supports image generation, and enables GradCAT interpretability.

Architecture: NesT processes non-overlapping image blocks with canonical transformers and uses hierarchical block aggregation for cross-block communication.The design requires only minor code changes to the original ViT.
Architecture: Hierarchically nested transformers with block aggregation outperform previous sophisticated local self-attention variants while simplifying the architecture.The aggregation provides a perspective for effective cross-block communication.
Results: 83.3%/83.8% ImageNet accuracy is obtained by NesT models with 38M/68M parameters.NesT also reaches 75.9%/82.3% training accuracy after 30/100 epochs, indicating fast convergence.
Results: 96% CIFAR10 accuracy is achieved by a 6M-parameter NesT trained using a single GPU.This example illustrates data-efficient learning on a small dataset.
Applications: Extending NesT to image generation produces a strong decoder with better performance than convolutional architectures and comparable speed on 64 × 64 ImageNet generation.The paper identifies efficient generative modeling as an application of this extension.
Interpretability: GradCAT interprets NesT by traversing its tree-like structure to explain how aggregated local transformers process semantic image patches.The method leverages decoupled feature learning and abstraction.

Related Work

Vision transformers have shown strong performance but remain difficult to train efficiently with limited data. Existing remedies add supervision or architectural complexity, while reported data-efficient performance can remain below convolutional networks.

Vision transformers: ViT models can outperform state-of-the-art convolutional networks when trained on tens or hundreds of millions of labeled samples.This establishes the large-data setting in which the original ViT performs strongly.
Data efficiency: Data efficiency is a central challenge because the original ViT requires hundreds of millions of images for pre-training.Such a data requirement is not always practical.
Data efficiency: DeiT addresses data requirements through convolutional-network teacher distillation, but this increases supervised training complexity.Reported performance on data-efficient benchmarks still significantly underperforms convolutional networks.

Proposed Method

NesT applies standard transformer blocks independently within non-overlapping image blocks and nests them through spatial aggregation. Aggregation on the full image plane exchanges information across blocks while progressively expanding receptive fields.

Input and hierarchy: NesT embeds image patches, partitions them into blocks, and processes every hierarchy with shared transformer parameters.The hierarchy is controlled by patch size S × S and the number of block hierarchies.
Interpretability: GradCAT uses NesT’s tree-like hierarchy to traverse class-related activations and locate image patches associated with target classes.Each tree node represents an averaged activation value used along the decision path.
Local processing: Canonical transformer layers combine multi-head self-attention, feed-forward networks, skip connections, and layer normalization within each block.Trainable positional embeddings encode spatial information before block processing.
Block aggregation: Every four spatially connected blocks merge into one, reducing the number of blocks by four while preserving sequence length.This creates a nested hierarchy with gradually expanding receptive fields.
Block aggregation: Block aggregation unblocks features to the full image plane, applies spatial operations, and blocks the downsampled maps for the next hierarchy.This is how independently processed blocks exchange and mix global information.
Block aggregation: The aggregation uses a 3 × 3 convolution, layer normalization, and 3 × 3 max pooling, with small kernels exchanging information across block boundaries.Ablations emphasize image-plane processing and applying convolution before pooling.
Design rationale: Unlike more complex local-attention methods, NesT uses standard independent block processing and communicates globally only during simple spatial aggregation.This targets cross-block communication without adding complexity to self-attention layers.

Transposed NesT for Image Generation

Transposing NesT yields a generative decoder by replacing block aggregation with block de-aggregation, enabling progressively more blocks and full-sized image output. The design achieves better generation quality and substantially faster decoding than prior transformer-based generation.

Results: The transposed NesT decoder has better performance than comparable-speed convolutional decoders and nearly an order-of-magnitude faster decoding than TransGAN.
Decoder design: NesT becomes a generator by replacing block aggregation with block de-aggregation, which progressively increases the number of blocks before producing a full-sized image.The input is a noise vector, and pixel shuffle upsamples feature maps during de-aggregation.
Decoder design: The decoder’s block count increases by a factor of 4 across hierarchies, ending with an output sequence that is unblocked into an H × W × 3 image.

GradCAT: Interpretability via Tree Traversal

GradCAT interprets NesT by tracing class-relevant information through its nested, tree-like hierarchy. It combines activations and class-specific gradients to select high-value paths from the root toward image-block leaves.

Motivation: NesT’s independent block processing and nested hierarchy resemble a decision tree, with aggregation selecting among non-overlapping local features.
GradCAT method: GradCAT traces the traversal path that contributes most to the predicted class logit, recursively following high-value information from the root to a leaf.
GradCAT method: At the top hierarchy, each child node processes one of four non-overlapping 2×2 feature-map partitions, providing the traversal’s spatial branching structure.
GradCAT method: The negative class-specific gradient indicates the ascent direction for increasing class-c logit, so higher positive values mark greater importance.

Experiments

Experiments evaluate NesT across classification, interpretability, and image generation, showing data-efficient learning, competitive accuracy, interpretable attention, and faster generative modeling. Ablations identify block aggregation and de-aggregation choices that affect performance.

Experimental setup: NesT is evaluated on ImageNet, CIFAR, interpretability, and 64 × 64 ImageNet generation using capacity-matched architectures.The experiments compare tiny, small, and base configurations and include classification, visualization, and generative-modeling settings.
Visual Interpretability: GradGAT traverses NesT’s tree structure using mean activation strengths and can locate predicted objects while exposing decision-path errors.The lighter example shows a correct prediction based on salient cues from a red light rather than the ground-truth matchstick.
Visual Interpretability: NesT CAM produces finer object-region attention than DeiT Rollout and less noise than ResNet50 GradCAM++.All compared models are trained on ImageNet.
Ablation Studies: NesT-B loses only 1.5% when reducing ImageNet training from 300 to 100 epochs, and shows less reliance on data augmentation.The convergence and augmentation results are shown in the cited ablation studies.
Ablation Studies: Block aggregation works best when spatial operations are applied on the holistic image plane, with small convolution before max pooling.The ablation reports max pooling as better than stride-2 subsampling and average pooling.

Conclusion

The paper concludes that nested-transformer aggregation improves data efficiency and convergence while extending to faster image generation and more interpretable feature extraction.

Conclusion: Nested-transformer aggregation matches more complex methods while improving data efficiency and convergence speed.The paper also reports speed gains for image generation and improved interpretability through gradient-based tree traversal.

Appendix

The appendix provides additional experimental results that supplement the paper’s main experiments.

Appendix: Additional experimental results are provided to complete the experimental sections of the main paper.

NesT Architecture and Training Details

NesT’s architecture study focuses on nested-transformer aggregation rather than per-block hyperparameters, using configurations adapted to image size and model scale.

Architecture configuration: The architecture analysis prioritizes nested-transformer aggregation and follows prior work for per-block hyperparameters and comparable model capacity.Capacity is matched using factors such as parameter count and throughput.
Architecture configuration: Patch size S × S and hierarchy depth Td determine the overall hierarchy, with configurations adjusted for small and large input sizes.Tiny, small, and base variants follow prior settings for head counts and hidden dimensions.

NesT Hierarchy Variants

NesT’s hierarchy depth and word size shape how narrowly transformers process image information. Experiments compare these configurations across model sizes and CIFAR datasets while documenting the training setup.

Hierarchy design: Deeper hierarchies make each block focus on a narrower range of pixel information, producing shorter sequence lengths.
Hierarchy design: Shallower NesT variants show a clear accuracy drop despite having the same total number of self-attention layers.At depth 1, NesT degenerates into a global self-attention method such as ViT.
Hierarchy design: Depth 5 yields a marginal decrease because each block has sequence length n = 2 × 2.
Experimental configurations: The experiments evaluate S = {1, 2} with depths {2, 3, 4, 5} on NesT-T, NesT-S, and NesT-B across two CIFAR datasets.
Experimental configurations: NesT4 with S = 1 is the default configuration for most CIFAR experiments.
Training setup: CIFAR models use AdamW, weight decay 0.05, a 0.1 stochastic-depth drop rate, and 5 warm-up epochs.The base learning rate is 2.5e-6 per device, with linear scaling by total batch size divided by 256.

Impact of Block Aggregation

The ablations examine aggregation, augmentation, architectural configurations, generation, and GradCAT interpretation. They identify non-trivial aggregation effects, stable augmentation behavior, strong visual interpretation results, and an object-size limitation.

Block aggregation: Block aggregation choices have a non-trivial impact on performance, although many options differ only slightly.Query down-sampling inside self-attention also makes transformers more difficult to train because skip connections require corresponding down-sampling.
Data augmentation: NesT remains highly stable when augmentation components are removed individually or consecutively, compared with DeiT.The authors speculate that effective vision cues are easier to learn with local than global attention and note comparable stability in Swin Transformer.
Teacher distillation: Weak teacher distillation raises NesT-B to 84.9% CIFAR100 accuracy after 300 epochs and 86.1% after 1000 epochs using 2 GPUs.
Teacher distillation: Strong augmentation for both teacher and target model decreases the effectiveness of teacher distillation, for both small and large teachers.The authors leave ImageNet verification for future work.
Image generation: NesT can be repurposed as a decoder for image generation by replacing block aggregation with block de-aggregation and using pixel shuffle.The hidden dimension is reduced to one quarter of the input in the described de-aggregation operation.
Visual interpretation: GradCAT visualizations are less meaningful when the target object is large and centered, a pattern reported for most ImageNet images.The authors identify datasets with non-centered objects as a direction for future study.
Visual interpretation: NesT CAM outperforms the compared methods and baselines for bounding-box generation while requiring only a single forward pass.The comparison searches each method’s best threshold on the ImageNet 50k validation set.
Block aggregation: Applying block aggregation on the image plane enables information exchange across block boundaries and among blocks entering different merged blocks.