Scaling White-Box Transformers for Vision

TL;DR

This paper introduces CRATE-$\alpha$, an improved scalable white-box transformer architecture for vision tasks, achieving higher accuracy on ImageNet while maintaining interpretability, through minimal modifications and a light training recipe.

Contribution

The paper presents CRATE-$\alpha$, a scalable version of the CRATE architecture with enhanced performance and interpretability for vision transformers, addressing previous scalability limitations.

Findings

CRATE-$\alpha$-B achieves 83.2% ImageNet accuracy, surpassing prior models.

CRATE-$\alpha$-L reaches 85.1% accuracy on ImageNet.

Model interpretability is preserved and improved with larger CRATE-$\alpha$ models.

Abstract

CRATE, a white-box transformer architecture designed to learn compressed and sparse representations, offers an intriguing alternative to standard vision transformers (ViTs) due to its inherent mathematical interpretability. Despite extensive investigations into the scaling behaviors of language and vision transformers, the scalability of CRATE remains an open question which this paper aims to address. Specifically, we propose CRATE-$\alpha$, featuring strategic yet minimal modifications to the sparse coding block in the CRATE architecture design, and a light training recipe designed to improve the scalability of CRATE. Through extensive experiments, we demonstrate that CRATE-$\alpha$ can effectively scale with larger model sizes and datasets. For example, our CRATE-$\alpha$-B substantially outperforms the prior best CRATE-B model accuracy on ImageNet classification by 3.7%, achieving an accuracy of 83.2%. Meanwhile, when scaling further, our CRATE-$\alpha$-L obtains an ImageNet classification accuracy of 85.1%. More notably, these model performance improvements are achieved while preserving, and potentially even enhancing the interpretability of learned CRATE models, as we demonstrate through showing that the learned token representations of increasingly larger trained CRATE-$\alpha$ models yield increasingly higher-quality unsupervised object segmentation of images. The project page is https://rayjryang.github.io/CRATE-alpha/.