Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding
This work addresses the need for simpler, more efficient vision transformers in computer vision, offering incremental improvements in data efficiency and interpretability.
The paper tackles the challenge of making vision transformers more data-efficient and interpretable by proposing a nested hierarchical architecture that simplifies design and reduces training data requirements, achieving faster convergence on ImageNet and CIFAR and an 8x speedup in image generation.
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.