Hierarchical Vision Transformers as Masked Autoencoders

In recent years, Vision Transformers (ViT) have garnered significant attention for their impressive performance in computer vision tasks. It challenged the traditional vision paradigm that relies on convolutions. The sequence-to-sequence paradigm is less intuitive for vision tasks compared to NLPs, but some architectures have proven to be undeniably strong at many vision tasks.

The Challenges of ViT with high-resolution images

As the resolution of images increases, the computational and memory demands of ViTs become increasingly prohibitive. Because the number of tokens grows quadratically with the image resolution! This quickly becomes prohibitive. This is where Hierarchical Vision Transformers (Hiera) shine as a more efficient and scalable alternative.

Vision Transformers rely on dividing an image into fixed-size non-overlapping patches, which are then treated as tokens for the model. The key issue arises from the fact that each token is processed through all stages of the transformer. As a result, when working with high-resolution images, the number of tokens increases dramatically, leading to a massive increase in computational complexity and memory usage. This makes ViT less feasible for large images, where processing all tokens at every stage becomes impractical.

How Hiera addresses these issues

Hiera, on the other hand, takes a hierarchical approach to token processing. Pooling the query tokens (Q) significantly improves efficiency. Unlike ViT, Hiera does not process all tokens at every stage. Instead, it selectively reduces the number of tokens passed through deeper layers. This hierarchical tokenization ensures that only the most relevant tokens are processed at higher levels, drastically reducing the memory footprint while maintaining or even improving model performance.

In high-resolution image tasks, this approach allows Hiera to maintain a more manageable number of tokens at each stage. By focusing on the most important information at each level, Hiera can capture fine-grained details without the need to process the entire set of tokens throughout the model. This makes Hiera a strong candidate for replacing ViT, especially when working with images of very high resolution.

Further development for higher resolution images

However, many groups, including mine at work, witnessed something unexpected. Hiera really doesn’t generalize well to different resolutions. This shouldn’t happen: the only component of a Vision Transformer that requires knowledge of the input resolution is the positional embedding. There is no reason for this to cause any problems for the higher-resolution downstream task. It turns out the problem arose from an interaction between Window Attention and the positional embedding. Hiera was then used in the trunk of the image encoder for SAM2. It was one of the key contributions that enabled the creation of this compelling real-time demo.

This was an exciting innovation to see, as it led to great improvements. In a field where we need more players, we need more papers that make it possible for smaller groups to propose SOTA systems with fewer resources.

In my experience. Pre-training, training and inference.

Hiera holds much promise as it addresses the prohibitive memory usage of non-hierarchical Vision Transformers. However, they still rely on a sequence-to-sequence paradigm that is not ideal for vision. In my experience, it appears that the inductive biases inherent to CNNs make for better models. Typically, ConvNeXt is a very hard model to beat in downstream prediction tasks.

In my experience, Hiera undeniably beats ViT. But it is still underwhelming in many aspects. First, the MAE pre-training step is computationally expensive. This is a heavy constraint on relatively GPU-poor groups. Second, the performance of the fine-tuning step is lackluster. Not only does it still require a decent amount of data for the prediction task, but the model also struggles to compete with its more advanced counterparts (ConvNeXt, PvT, MViT, Swin-V2) that are initialized from ImageNet and trained directly on the task (with no specific pre-training task). I haven’t managed to make it really competitive despite my best efforts and extensive experimentation (different masking ratios, layer-wise learning rate decay, high batch sizes, and many different fine-tuning settings). It gets better if you keep the earlier layers unchanged (prefer downsampling the image over cropping, keep the patch size fixed between pre-training and fine-tuning, etc.), but it never becomes SOTA on the proprietary datasets I worked on. Third, it looks like Hiera has some convergence issues. The pre-training step might result in occasional crashes where the loss suddenly becomes NaN and all the weights are messed up. All of this is for a speed increase in the forward and backward passes that is marginal for relatively high-resolution images (~3000 × 2000).

Closing thoughts

I think Hiera and MAE in general are interesting approaches. But they offer little utility when you try to apply them without some of the biggest datasets out there, coupled with serious GPU power. Nonetheless, I think pretext tasks for pre-training make sense and might help in many downstream scenarios. I’m excited to read more research on this topic.