Scaling White-Box Transformers for Vision

Authors: jinrui  yang
xianhang  li
druv  pai
yuyin  zhou
yi  ma
yaodong  yu
cihang  xie
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 00:21:51

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.

Main paper: 
 introduction: 
 Over the past several years, the Transformer architecture [42] has dominated deep representation learning for natural language processing (NLP), image processing, and visual computing [8, 2, 9, 5, 12] . However, the design of the Transformer architecture and its many variants remains largely empirical and lacks a rigorous mathematical interpretation. This has largely hindered the development of new Transformer variants with improved efficiency or interpretability. The recent white-box Transformer model CRATE [46] addresses this gap by deriving a simplified Transformer block via unrolled optimization on the so-called sparse rate reduction representation learning objective.More specifically, layers of the white-box CRATE architecture are mathematically derived and fully explainable as unrolled gradient descent-like iterations for optimizing the sparse rate reduction. The self-attention blocks of CRATE explicitly conduct compression via denoising features against learned low-dimensional subspaces, and the MLP block is replaced by an incremental sparsification (via ISTA [1, 11] ) of the features. As shown in previous work [47] , besides mathematical interpretability, the learned CRATE models and features also have much better semantic interpretability than conventional transformers, i.e., visualizing features of an image naturally forms a zero-shot image segmentation of that image, even when the model is only trained on classification.Scaling model size is widely regarded as a pathway to improved performance and emergent properties [44, 40, 41, 14] . Until now, the deployment of CRATE has been limited to relatively modest scales. The most extensive model described to date is the base model size encompasses 77.6M parameters 38th Conference on Neural Information Processing Systems (NeurIPS 2024). We compare the FLOPs and accuracy on ImageNet-1K of our methods with ViT [9] and CRATE [46] . The values of CRATE-α model correspond to those presented in Table 1 . A more detailed comparison between CRATE-α and ViT is included in Appendix A.2.(CRATE-Large) [46] . This contrasts sharply with standard Vision Transformers (ViTs [9] ), which have been effectively scaled to a much larger model size, namely 22B parameters [5] .To this end, this paper provides the first exploration of training CRATE at different scales for vision, i.e., Tiny, Small, Base, Large, Huge. Detailed model specifications are given in Table 7 of Appendix A.1. To achieve effective scaling, we make two key changes. First, we identify the vanilla ISTA block within CRATE as a limiting factor that hinders further scaling. To overcome this, we significantly expand the channels, decouple the association matrix, and add a residual connection, resulting in a new model variant -CRATE-α. It is worth noting that this architecture change still preserves the mathematical interpretability of the model. Second, we propose an improved training recipe, inspired by previous work [38, 46, 39] , for better coping the training with our new CRATE-α architecture.We provide extensive experiments supporting the effective scaling of our CRATE-α models. For example, we scale the CRATE-α model from Base to Large size for supervised image classification on ImageNet-21K [6] , achieving 85.1% top-1 accuracy on ImageNet-1K at the Large model size.We further scale the model size from Large to Huge, utilizing vision-language pre-training with contrastive learning on DataComp1B [10] , and achieve a zero-shot top-1 accuracy of 72.3% on ImageNet-1K at the Huge model size. 1 These results demonstrate the strong scalability of the CRATE-α model, shedding light on scaling up mathematically interpretable models for future work.The main contributions of this paper are threefold: 1. We design three strategic yet minimal modifications for the CRATE model architecture to unleash its potential. In Figure White-box Transformers. [46, 45] argued that the quality of a learned representation can be assessed through a unified objective function called the sparse rate reduction. Based on this framework, [46, 45] developed a family of transformer-like deep network architectures, named CRATE, which are mathematically fully interpretable. CRATE models has been demonstrably effective on various tasks, including vision self-supervised learning and language modeling [26, 45] . Nevertheless, it remains unclear whether CRATE can scale as effectively as widely used black-box transformers. Previous work [46] suggests that scaling the vanilla CRATE model can be notably challenging.Scaling ViT. ViT [9] represents the initial successful applications of Transformers to the image domain on a large scale. Many works [12, 31, 33, 32, 5, 37, 21, 22, 29, 18, 49] have deeply explored various ways of scaling ViTs in terms of model size and data size. From the perspective of selfsupervision, MAE [12] provides a scalable approach to effectively training a ViT-Huge model using only ImageNet-1K. Following the idea of MAE, [31] further scales both model parameters to billions and data size to billions of images. Additionally, CLIP was the first to successfully scale ViT on a larger data scale (i.e., 400M) using natural language supervision. Based on CLIP, [32, 33] further scale the model size to 18 billion parameters, named EVA-CLIP-18B, achieving consistent performance improvements with the scaling of ViT model size. From the perspective of supervised learning, [49, 5] present a comprehensive analysis of the empirical scaling laws for vision transformers on image classification tasks, sharing some similar conclusions with [15] . [49] suggests that the performancecompute frontier for ViT models, given sufficient training data, tends to follow a saturating power law. More recently, [5] scales up ViT to 22 billion parameters. Scaling up different model architectures is non-trivial. [37, 21, 22] have made many efforts to effectively scale up different architectures. In this paper, due to the lack of study on the scalability of white-box models, we explore key architectural modifications to effectively scale up white-box transformers in the image domain. 
background and preliminaries: 
 In this section, we present the background on white-box transformers proposed in [46] , including representation learning objectives, unrolled optimization, and model architecture. We first introduce the notation that will be used in the later presentation.Notation. We use notation and problem setup following Yu et al. [46] . We use the matrix-valued random variable X = [x 1 , . . . , x N ] ∈ R D×N to represent the data, where each x i ∈ R D is a "token", such that each data point is a realization of X. For instance, X can represent a collection of image patches for an image, and x i is the i-th image patch. We use f ∈ F : R D×N → R d×N to denote the mapping induced by the transformer, and we let Z = f (X) = [z 1 , . . . , z N ] ∈ R d×N denote the features for input data X. Specifically, z i ∈ R d denotes the feature of the i-th input token x i . The transformer f consists of multiple, say L, layers, and so can be written asf = f L • • • • • f 1 • f pre ,where f ℓ : R d×N → R d×N denotes the ℓ-th layer of the transformer, and the pre-processing layer is denoted by f pre = R D×N → R d×N . The input to the ℓ-th layer f ℓ of the transformer is denoted byZ ℓ = z ℓ 1 , . . . , z ℓ N ∈ R d×N , so that f ℓ : Z ℓ → Z ℓ+1 .In particular, Z 1 = f pre (X) ∈ R d×N denotes the output of the pre-processing layer and the input to the first layer.Following the framework proposed in [45] , we posit that the goal of representation learning is to learn a feature mapping or representation f ∈ F : R D×N → R d×N that transforms the input data X (which may have a nonlinear, multi-modal, and otherwise complicated distribution) into structured and compact features Z, such that the token features lie on a union of low-dimensional subspaces, say with orthonormal bases [46] proposes the Sparse Rate Reduction (SRR) objective to measure the goodness of such a learned representation:U [K] = (U k ) k∈[K] ∈ (R d×p ) K .EQUATIONwhere Z = f (X) denotes the token representation, ∥Z∥ 1 denotes the ℓ 1 norm, and R(Z), R c (Z | U [K] ) are (estimates for) rate distortions [4, 7] , defined as:EQUATIONIn particular, R c (Z | U [K] ) (resp. R(Z)) provide closed-form estimates for the number of bits required to encode the sample Z up to precision ϵ > 0, conditioned (resp. unconditioned) on the samples being drawn from the subspaces with bases U [K] . Minimizing the term R c improves the compression of the features Z against the posited model, and maximizing the term R promotes non-collapsed features. The remaining term λ∥Z∥ 1 promotes sparse features. Refer to [45] for more details about the desiderata and objective of representation learning via the rate reduction approach.Unrolled optimization. To optimize the learning objective and learn compact and structured representation, one approach is unrolled optimization [11, 36] : each layer of the deep network implements an iteration of an optimization algorithm on the learning objective. For example, one can design the layer f ℓ such that the forward pass is equivalent to a proximal gradient descent step for optimizing learning objective L(Z), i.e.,Z ℓ+1 = f ℓ (Z ℓ ) = Prox[Z ℓ -η • ∇ Z L(Z ℓ )].Here we use η to denote the step size and Prox[•] to denote the proximal operator [27] .One layer of the CRATE model. We now present the design of each layer of the white-box transformer architecture -Coding RATE Transformer (CRATE) -proposed in [46] . Each layer of CRATE contains two blocks: the compression block and the sparsification block. These correspond to a two-step alternating optimization procedure for optimizing the sparse rate reduction objective (1) . Specifically, the ℓ-th layer of CRATE is defined asEQUATIONCompression block (MSSA). The compression block in CRATE, called Multi-head Subspace Self-Attention block (MSSA), is derived for compressing the token set Z = [z 1 , . . . , z N ] by optimizing the compression term R c (defined Eq. ( 1)), i.e.,EQUATIONwhere U ℓ [K] denotes the (local) signal model at layer ℓ, and the MSSA operator is defined asEQUATIONCompared with the commonly used attention block in transformer [42] , where the k-th attention head is defined as(V ⊤ k Z) softmax((Q ⊤ k Z) ⊤ (K ⊤ k Z)), MSSA uses only one matrix to obtain the query, key, and value matrices in the attention: that is,U k = Q k = K k = V k .Sparse coding block (ISTA). The Iterative Shrinkage-Thresholding Algorithm (ISTA) block is designed to optimize the sparsity term and the global coding rate term, 1). [46] shows that an optimization strategy for these terms posits a (complete) incoherent dictionary D ℓ ∈ R d×d and takes a proximal gradient descent step towards solving the associated LASSO problemλ∥Z∥ 0 -R(Z | U [K] ) in (arg min Z≥0 [ 1 2 ∥Z ℓ+1/2 -D ℓ Z∥ 2 2 + λ∥Z∥ 1 ], obtaining the iterationEQUATIONIn particular, the ISTA block sparsifies the intermediate iterates Z ℓ+1/2 w.r.t. D ℓ to obtain Z ℓ+1 . 
crate-α model: 
 In this section, we present the CRATE-α architecture, which is a variant of CRATE [46] . As shown in Fig. 1 (Right), there is a significant performance gap between the white-box transformer CRATE-B/16ISTA-overcomplete MSSA ODL ODL MSSA ODL • • • • • • Aggregate Figure 2:One layer of the CRATE-α model architecture. MSSA (Multi-head Subspace Self-Attention, defined in ( 5)) represents the compression block, and ODL (Overcomplete Dictionary Learning, defined in ( 12)) represents the sparse coding block. A more detailed illustration of the modifications is provided in Fig. 6 in the Appendix .(70.8%) and the vision transformer ViT-B/16 (84.0%) [9] . One possible reason is that the ISTA block applies a complete dictionary D ∈ R d×d , which may limit its expressiveness. In contrast, the MLP block in the transformer2 applies two linear transformations W 1 , W 2 ∈ R d×4d , leading to the MLP block having 8 times more parameters than the ISTA block.Since the ISTA block in CRATE applies a single incremental step to optimize the sparsity objective, applying an orthogonal dictionary can make it ineffective in sparsifying the token representations.Previous work [28] has theoretically demonstrated that overcomplete dictionary learning enjoys a favorable optimization landscape. In this work, we use an overcomplete dictionary in the sparse coding block to promote sparsity in the features. Specifically, instead of using a complete dictionary D ℓ ∈ R d×d , we use an overcomplete dictionary D ℓ ∈ R d× (Cd) , where C > 1 (a positive integer) is the overcompleteness parameter. Furthermore, we explore two additional modifications to the sparse coding block that lead to improved performance for CRATE. We now describe the three variants of the sparse coding block that we use in this paper.Modification #1: Overparameterized sparse coding block. For the output of the ℓ-th CRATE attention block Z ℓ+1/2 , we propose to sparsify the token representations with respect to an overcomplete dictionary D ℓ ∈ R d×(Cd) by optimizing the following LASSO problem,EQUATIONTo approximately solve (7) , we apply two steps of proximal gradient descent, i.e.,EQUATION)where Prox is the proximal operator of the above non-negative LASSO problem (7) and defined asEQUATIONThe output of the sparse coding block is defined asEQUATIONNamely, A ℓ is a sparse representation of Z ℓ+1/2 with respect to the overcomplete dictionary D ℓ . The original CRATE ISTA tries to learn a complete dictionary D ∈ R d×d to transform and sparsify the features Z. By leveraging more atoms than the ambient dimension, the overcomplete dictionary D ∈ R d×(Cd) can provide a redundant yet expressive codebook to identify the salient sparse structures underlying Z. As shown in Fig. 1 , the overcomplete dictionary design leads to 5.3% improvement compared to the vanilla CRATE model.Modification #2: Decoupled dictionary. We propose to apply a decoupled dictionary D ℓ in the last step (defined in (10) of the sparse coding block, Z ℓ+1 = D ℓ A ℓ , where D ℓ ∈ R d×(Cd) is a different dictionary compared to D ℓ . By introducing the decoupled dictionary, we further improve the model performance by 2.0%, as shown in Fig. 1 . We denote this mapping from Z ℓ+1/2 to Z ℓ+1 as the Overcomplete Dictionary Learning block (ODL), defined as follows:EQUATION)Modification #3: Residual connection. Based on the previous two modifications, we further add a residual connection, obtaining the following modified sparse coding block:EQUATIONAn intuitive interpretation of this modified sparse coding block is as follows: instead of directly sparsifying the feature representations Z, we first identify the potential sparse patterns present in Z by encoding it over a learned dictionary. Subsequently, we incrementally refine Z by exploiting the sparse codes obtained from the previous encoding step. From Fig. 1 , we find that the residual connection leads to a 0.7% improvement.To summarize, to effectively scale white-box transformers, we implement three modifications to the vanilla white-box CRATE model proposed in [46] . Specifically, in our CRATE-α model, we introduce a decoupling mechanism, quadruple the dimension of the dictionary (4×), and incorporate a residual connection in the sparse coding block. 
experiments: 
 Overall. The experimental section consists of three parts: (1) Scaling study: We thoroughly investigate the scaling behaviors of CRATE-α from Base to Large size and ultimately to Huge size. ( 2) Downstream applications: To further verify the broader benefits of scaling the CRATE-α model, we conduct additional experiments on real-world downstream tasks and present preliminary exploration results of CRATE-α on language tasks. (3) Interpretability: In addition to scalability, we study the interpretability of CRATE-α across different model sizes.Scaling Study. For the transition from Base to Large size, we pre-train our model on ImageNet-21K and fine-tune it on ImageNet-1K via supervised learning. When scaling from Large to Huge, we utilize the DataComp1B [10] dataset within a vision-language pre-training paradigm, allowing us to study the effects of scaling the model to a massive size. For evaluation, we evaluate the zero-shot accuracy of these models on ImageNet-1K. Downstream Applications. We include additional experimental results on four downstream datasets (CIFAR-10/100, Oxford Flowers, and Oxford-IIT Pets). We also examine the dense prediction capability of CRATE-α by training it on segmentation tasks using the ADE20K dataset [51] . For language tasks, we conduct new experiments with CRATE-α using autoregressive training on OpenWebText, following the setup in nanoGPT [16] .Interpretability. Following the evaluation setup of CRATE as outlined in [46] , we apply MaskCut [43] to validate and evaluate the rich semantic information captured by our model in a zero-shot setting, including both qualitative and quantitative results. , weight decay to 0.1, and batch size to 4096. We apply data augmentation techniques such as Inception crop [35] resized to 224 and random horizontal flipping. In the fine-tuning phase, we adjust the base learning rate to 1.6 × 10 -4 , maintain weight decay at 0.1, and batch size at 4096. We apply label smoothing with a smoothing parameter of 0.1 and apply data augmentation methods including Inception crop, random horizontal flipping, and random augmentation with two transformations (magnitude of 9). For evaluation, we resize the smaller side of an image to 256 while maintaining the original aspect ratio and then crop the central portion to 224×224. In both the pre-training and fine-tuning phases, we use the AdamW optimizer [24] and incorporate a warm-up strategy, characterized by a linear increase over 10 epochs. Both the pre-training and fine-tuning are conducted for a total of 91 epochs, utilizing a cosine decay schedule.(2) Large to Huge size: In the pre-training stage, we utilize an image size of 84×84, and the maximum token length is 32, with a total of 2.56 billion training samples. During the fine-tuning stage, we increase the image size to 224×224 while maintaining the maximum token length at 32, with a 512 million training samples. Here, the key distinction between the pre-training stage and the fine-tuning stage is the image size. A smaller image size results in a faster training speed. In the configurations of CRATE-α-CLIPA-B, CRATE-α-CLIPA-L, and CRATE-α-CLIPA-H, we use the CRATE-α model as the vision encoder, and utilize the same pre-trained huge transformer model from CLIPA [18] as the text encoder. For both the pre-training and fine-tuning stages, we freeze the text encoder and only train the vision encoder, i.e., the CRATE-α model. As we will show in the later results, this setup effectively demonstrates the scaling behaviors of CRATE-α models in the image domain. Detailed hyperparameter settings can be found in Appendix A.Downstream Applications. On four downstream datasets, we follow the training setup from [46] . For the segmentation task, we compare the performance of CRATE and CRATE-α on the ADE20K dataset, mainly following the setup of [30] with minor modifications. Our batch size is set to 128, and the total number of training steps is 5000. For the language task, we conduct experiments with CRATE-α using autoregressive training on OpenWebText, following the setup in [16] . We compare CRATE-α models (57M and 120M) with CRATE and GPT-2, using results from CRATE reported in [45] . Pre-train Fine-tune Scaling the CRATE-α Model from Base to Large. As shown in Table 1 , we compare CRATE-α-B and CRATE-α-L at patch sizes 32, 16, and 8. Firstly, we find our proposed CRATE-α-L consistently achieves significant improvements across all patch sizes. Secondly, compared with the results of the vanilla CRATE (the first row of Table 1 ), increasing from CRATE-B to CRATE-L results in only a 0.5% improvement on ImageNet-1K. This indicates a case of diminishing returns. These findings compellingly highlight that the scalability of CRATE-α models significantly outperforms that of the vanilla CRATE. Meanwhile, the training loss in the pre-training stage is presented in Fig. 3 ; as the model capacity increases, the trend of the training loss improves predictably. This phenomenon is also described in [9] .Scaling the CRATE-α Model from Large to Huge. From the results shown in Fig. 4 , we find that:(1) CRATE-α-CLIPA-L/14 significantly outperforms CRATE-α-CLIPA-B/16 by 11.3% and 9.0% in terms of ImageNet-1K zero-shot accuracy during the pre-training and fine-tuning stages, respectively. The substantial benefit suggests that the quality of learned representation may be constrained by the model size. Therefore, increasing the model size effectively leverages larger amounts of data. (2) When continuing to scale up model size, we also observe that CRATE-α-CLIP-H/14 continues to benefit from larger training datasets, outperforming CRATE-α-CLIP-L/14 by 3.1% and 2.5% in terms of ImageNet-1K zero-shot accuracy during the pre-training and fine-tuning stages, respectively. This demonstrates the strong scalability of the CRATE-α model. To explore the performance ceiling, we train a standard ViT-CLIPA-H/14 from scratch and observe improved performance. Downstream Applications. On four downstream datasets, as shown in Table 2 , we find that CRATE-α consistently outperforms CRATE, with both models pre-trained on IN21K, while CRATE-α demonstrates improved performance as model size increases. For the segmentation task, results in Table 3 show that CRATE-α consistently outperforms CRATE across all key metrics, with both models pre-trained on IN21K. These findings indicate significant performance gains in vision tasks beyond classification. For the language task, Table 4 shows that CRATE-α significantly improves over CRATE in language modeling. Due to limited time and resource constraints, we completed 80% of the total iterations for CRATE-α-small and 55% for CRATE-α-base, compared to the 600K total iterations used for CRATE. Nevertheless, CRATE-α still demonstrated notable improvements.Interpretability. As shown in Fig. 5 , we provide the segmentation visualization on COCO val2017 [20] for CRATE-α, CRATE, and ViT, respectively. We find that our model preserves and even improves the (semantic) interpretability advantages of CRATE. Moreover, we summarize quantitative evaluation results on COCO val2017 in Table 6 . Interestingly, when scaling up model size for CRATE-α, the Large model improves over the Base model in terms of object detection and segmentation.We further explore methods to scale models efficiently in terms of computation. Table 1 demonstrates that the CRATE-α model scales effectively from the Base model to its larger variants. However, the pre-training computation for the top-performing model, CRATE-α-L/8, is resource-intensive on ImageNet-21K. Inspired by CLIPA [18] , we aim to reduce computational demands by using reduced image token sequence lengths, while maintaining the same training setup during the fine-tuning stage.The results are summarized in Table 5 .(1) When fine-tuning with CRATE-α-L/14 and using CRATE-α-L/32 for pretraining on ImageNet-21K, this approach consumes about 35% of the TPU v3 core-hours required by CRATE-α-L/14, yet achieves a promising 83.7% top-1 accuracy on ImageNet-1K, comparable to the 83.9% achieved by CRATE-α-L/14; (2) When fine-tuning with CRATE-α-L/8 and using CRATE-α-L/32 for pre-training, this approach consumes just 15% of the training time required by CRATE-α-L/8, yet it still achieves a promising 84.2% top-1 accuracy on ImageNet-1K, compared to 85.1% when using the CRATE-α-L/8 model in the pre-training stage; (3) While the total computational cost of CRATE-α-L/32 + CRATE-α-L/8 is less than that of CRATE-α-L/14 + CRATE-α-L/14, the performance of the former is slightly better. In summary, we find that this strategy offers a valuable reference for efficiently scaling CRATE-α models in the future. 
discussion: 
 Limitations. Although we have used some existing compute-efficient training methods (e.g., CLIPA [18] ) and have initiated an exploration into compute-efficient scaling strategies for whitebox transformers in Section 4.4, this work still requires a relatively large amount of computational resources, which may not be easily accessible to many researchers.Societal impact. A possible broader implication of this research is the energy consumption needed to conduct the experiments in our scaling study. However, there is growing interest in developing white-box transformers for better interpretability and transparency across a wide range of tasks and domains, including image segmentation [46] , self-supervised masked autoencoders [26] , and integrated sensing and communications [50] , etc. Moreover, our results on the scalability of white- Table 6 : Object detection and fine-grained segmentation via MaskCut on COCO val2017 [20] .We evaluate models of various scales and assess their average precision using COCO's official evaluation metric. Compared with existing models such as CRATE and ViT, CRATE-α model achieves a notable performance gain. In addition, when scaling CRATE-α from base to large, it also exhibits the benefit of scalability. box transformers could also shed light on scaling up a broader class of white-box deep neural networks, such as white-box ISTA networks and their variants [11, 34, 3, 48, 17] , designed via unrolled optimization. In summary, we believe that our findings and insights could be helpful for developing white-box transformers for a wide range of applications and tasks, benefiting a broad audience interested in building more interpretable and performant deep learning models and further amortizing the pre-training compute costs.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [No] Justification: We use public datasets and will also release the code. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 This paper provides the first exploration of training white-box transformer CRATE at scale for vision tasks. We introduce both principled architectural changes and improved training recipes to unleash the potential scalability of the CRATE type architectures. With these modifications, we successfully scale up the CRATE-α model along both the dimensions of model size and data size, while preserving, in most cases even improving, the semantic interpretability of the learned white-box transformer models. We believe this work provides valuable insights into scaling up mathematically interpretable deep neural networks, not limited to transformer-like architectures.with the CLIPA setup, warm-up steps are set to Visualization of self-attention maps of CRATE-α. We provide visualization of attention maps of CRATE-α in Fig. 7 . • The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide details of reproducing the experimental results in Section 4 and the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We evaluate our method on a large public validation dataset, which is representative.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We provide details about the training compute resources in Section 4.Guidelines:• The answer NA means that the paper does not include experiments. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Our work has no potential risk for misuse.Answer: [NA] Justification: This paper does not involve crowdsourcing or research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Model configurations are detailed in Table 7 (in Appendix A.1).The MLP block is defined asZ ℓ+1 = Z ℓ + W2σ(W ⊤ 1 Z ℓ+1/2), where σ is the nonlinear activation function and Z ℓ+1/2 denotes the output of the attention block. 


Decision: After rebuttal and discussion, all reviewers agree in their recommendation to accept this paper. The AC and SAC see no reason to overturn this recommendation, and hence decide to accept this paper. Congratulations.

Reviews: review: paper_summary: paper_summary: This paper aims to train CRATE at a large scale for vision tasks. The contribution includes an architecture modification to the sparse coding block and a light training recipe. The new model, called CRATE-alpha, shows large improvements compared with the previous CRATE model. The experiments also show promising results on unsupervised object segmentation.
strength_weakness: strength_weakness: Overall I find the paper is well-presented and solid. Below are my minor concerns for this paper:
- The paper is highly centered on improving CRATE. Most of the findings might not be transferable to other models. This may limit its impact to the general audience in NuerIPS community.
- It would be interesting to further understand its potential downstream applications (not only vision but also language data)
questions: questions: see weakness
limitations: limitations: N/A
paper_summary: paper_summary: This paper studies the scalability problem of white-box transformer CRATE and proposes CRATE-$\alpha$ to enhance the scaling ability of CRATE. To be specific, the authors propose three strategic but minimal modifications for the CRATE model architecture: Overparameterized sparse coding block, Decoupled dictionary, and Residual connection. Extensive experiments across different datasets and settings demonstrate the effectiveness of the proposed approach.
strength_weakness: strength_weakness: 1. Performance gaps with vanilla ViT. As shown in Figure 1, CRATE-$\alpha$ still lags behind vanilla ViT across different scales remarkably which may limit its application in real scenarios. Besides, it is suggested to compare with vanilla ViT in computational costs, number of parameters, and inference speed as well.

2. According to the model configuration, the number of parameters of CRATE-$\alpha$ is almost four times as CRATE and it is strange to consider those as the same scale models. Moreover, how do the proposed new modules contribute to the performance gain of CRATE-$\alpha$? Is it simply because of larger models?

3. Although the authors made lots of efforts in scaling CRATE to CRATE-$\alpha$, they only spent limited space in the paper to discuss the interpretability of the proposed method. This short paragraph may not be enough to justify why the authors are motivated to study the white-box transformers.
questions: questions: Apart from the questions in weakness above, another question is:

why the performance in dense prediction tasks is so bad?
limitations: limitations: N/A
paper_summary: paper_summary: This paper explores how to train white-box Transformers at scale for visual tasks. The authors propose a new model architecture called CRATE-$\alpha$, which extends the sparse coding block of the original CRATE model. A series of CRATE-$\alpha$ models were trained with varying model sizes, data sizes, and patch sizes using optimized training recipes. The main experiments focus on supervised classification and contrastive CLIP learning, with additional demonstrations of unsupervised semantic segmentation capability.
strength_weakness: strength_weakness: **Clarity:**
- The paper is heavily symbolized, relying extensively on intricate mathematical formulations rather than clear diagrams and straightforward language. Although this maintains academic rigor and professionalism, it severely hampers understanding of the paper's details and the broader dissemination of the model. Incorporating corresponding illustrations to explain the three modifications and comparing them with the standard Transformer structure would be beneficial.
- The organization of Section 4 is not concise, making it easy for readers to lose track.
  - The distinction between the paragraphs "Dataset and Evaluation" and "Training & Fine-tuning" is not well-defined, especially with the scattered descriptions of the data used.
  - The frequent interleaving of experimental setup descriptions with the presentation of experimental results disrupts the flow and coherence of the narrative.

**Significance:** 
- Although CRATE-$\alpha$ shows significant improvements over the original CRATE model, it still lags behind the state-of-the-art. For example, in the right side of Figure 1, CRATE-$\alpha$ typically requires nearly double the training FLOPs to achieve the same accuracy as ViT. 
- If the scalability and interpretability of a white-box Transformer architecture does not offer substantial insights and improvements, practitioners might prefer models with stronger performance but lower interpretability.
questions: questions: 1. As previously mentioned, as shown on the right side of Figure 1, CRATE-$\alpha$ usually requires approximately twice the FLOPs to reach the performance level of ViT, putting it at a noticeable disadvantage.
   
2. How does the performance improvement of CRATE-$\alpha$ compare to the original CRATE? Neither the CRATE models in Table 1 nor Figure 1 were pretrained on ImageNet-21K. Why was this not included for a fair comparison?
   
3. Lines 232-233 and Figure 3 describe the model’s **training loss** as predictable. Why not the **validation loss**, which is the primary concern in scaling laws and practical applications?
   
4. Table 2 only shows the compute requirements for the **pre-training stage**. Why does it not include the **fine-tuning** stage? Considering the total computational effort, I would like to see a comparison of the two scaling strategies: *CRATE-$\alpha$-L/32 + CRATE-$\alpha$-L/8* versus *CRATE-$\alpha$-L/14 + CRATE-$\alpha$-L/14*.
   
5. How was the amount of training data determined? Was there a specific standard or a FLOPs constraint? For example:
   - In Section 4.1, for training models from Base to Large, both pre-training and fine-tuning were conducted for a total of **91** epochs.
   - In Section 4.1, for training models from Large to Huge, there were **2.56** billion and **512** million training samples, respectively.
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: paper_summary: This paper introduces CRATE-α, an enhanced variant of the CRATE (Coding RATE Transformer) architecture, designed to scale efficiently while maintaining mathematical interpretability. The authors address the open question of CRATE's scalability by proposing strategic modifications to the sparse coding block and a refined training recipe. Extensive experiments demonstrate CRATE-α's effectiveness, showcasing improved performance on ImageNet classification tasks compared to the original CRATE model. Notably, the CRATE-α-B model achieved an 83.2% accuracy rate, a significant improvement over the previous best CRATE-B model.
strength_weakness: strength_weakness: Could the proposed architecture work well on other tasks like NLP?


While the paper provides a detailed analysis of the model's performance on ImageNet, there might be a need for more discussion on how these results generalize to other datasets and real-world applications.
questions: questions: See the weakness.
limitations: limitations: Yes

Comments: title: Response to Reviewer cPqJ
Discussion with Reviewer cPqJ (Part 3)
Discussion with Reviewer cPqJ (Part 2)
Discussion with Reviewer cPqJ (Part 1)
Response to Reviewer yemz
Post-rebuttal comment
comments
comment: Thank you for your valuable and in-depth review, as well as the series of insightful discussions. Your feedback and suggestions will undoubtedly enhance the quality of the paper. We greatly appreciate your input.

Regarding the 'more equitable experimental results', according to the NeurIPS discussion policy, "*The discussion phase is meant to clarify these questions, rather than to provide further comments regarding the reviews.*" Due to limited time and resource constraints, we are not able to provide new experimental results under more equitable settings before the end of the discussion phase. We will conduct new experiments as per your suggestions and incorporate them into our revised version.
Thank you for the author's prompt clarification; that makes things a bit clearer. However, my initial question and previous comments were specifically concerning Figure 1 (Right) and Table 1, and did not involve Figure 1 (Left).

The same questions remain: if they have different training processes and training data, this *"intuitive comparison of different published results"* weakens the support for the paper's conclusions. A fair comparison is highly valued in the field of computer vision.

I strongly recommend that the authors include more equitable experimental results in the paper, as this would greatly enhance its quality. If such changes are not feasible at this stage, I believe my rating is fair and reasonable.
Thank you for the detailed response. I appreciate the author's candid acknowledgment of the limitations of the current work. Given that these limitations objectively exist, this discussion will not change my evaluation.

I would like to clarify two points:
- The 7% I mentioned is intended to highlight the relative improvement potential, calculated as $(85.2 - 84.6) / 84.6 \times 100 = 7 $%.
- I hope the author is aware that, as shown in Figure 1 (Right), the additional FLOPs required for CRATE-α to match the comparable performance of ViT increases significantly with scale.
> **Q12**: Regarding A6, the authors' explanation **conflicts** with the paper. Please refer to the last two lines of the Figure 1 caption: *"CRATE is trained only on ImageNet-1K, while CRATE-α and ViT are pre-trained on ImageNet-21K."* Why do the authors now explain that *"all four model variants are pretrained on IN21K and then fine-tuned on IN1K"*? I kindly request the authors to consider my question closely: since CRATE has not been pre-trained on ImageNet-21K while CRATE-α has been *"pre-trained on ImageNet-21K and then fine-tuned on ImageNet-1K"* (as indicated in the caption of Table 1), how can these two models be compared fairly?

**A12:**  Thank you for your question, we apologize for the confusion by the unclear description in our submission. We would like to clarify that the comparison between CRATE and CRATE-α is fair. Specifically, the caption "CRATE is trained only on ImageNet-1K, while CRATE-α and ViT are pre-trained on ImageNet-21K" **applies only to Figure 1 (right), instead of both left and right figures**. 

In particular, we would like to clarify that the training setups for the CRATE model in Figure 1 (left) and Figure 1 (right) are different:

- The results for CRATE-B/16 and CRATE-L/16 in Figure 1 (right) and Table 1 are directly obtained from the original CRATE paper [YBP+2023], where the models were trained only on IN1K. This serves as an intuitive comparison of different published results. 
- The motivation for Figure 1 (left) is to ablate the effectiveness of our proposed three modifications. **It is an ablation study.** Therefore, we kept all training setups the same except for the model architecture. This is the reason — in **A6** of the rebuttal response — we mentioned that "all four model variants are pretrained on IN21K and then fine-tuned on IN1K." 

Additionally, according to the results of **Fine-tuning classification datasets** in **Q2: Additional experimental results on real-world downstream tasks** in the common response (we also applied the same training setup here), these results further validate the improvement of CRATE-α compared to CRATE under the same pre-trained dataset (IN21K). In the revision, we will clarify the different training/dataset setups and make sure to deliver clear messages to the readers. 

We again thank you for your review and valuable feedback, and hope we have provided satisfactory responses to your questions. Please let us know if you have further questions or comments.


[YBP+2023] White-Box Transformers via Sparse Rate Reduction. Yaodong Yu, Sam Buchanan, Druv Pai, et al. NeurIPS 2023.
>**Q11**: *Regarding A4, the authors reiterate that "it is very possible to scale up white-box models that are designed from principles to achieve the same performance as black-box ones." However, I would like to ask what practical benefits this white-box design (the so-called interpretability) has, aside from scaling up. For example, if this classification model performs poorly on a certain category, can this interpretability guide me in adjusting the model structure to improve performance in those categories? If the only advantage is scaling up, then such interpretability holds no real benefit, as black-box transformers can scale up more effectively (Figure 1 Right).*

**A11**: Thank you for your question regarding the practical benefits of this interpretable white-box design. Firstly, we would like to emphasize that model performance (such as accuracy) is not the only metric to evaluate a model. As we described in this work and shown in the prior work on white-box models [GL2010, CLW+2018, ZTA+2019, YBP+2023], we believe building understandings of how deep learning models work is also an important property for building trustworthy models. For example, 'advancing ongoing research in AI safety, including on the interpretability of AI systems' decision-making processes and on increasing the robustness of AI systems against misuse.' from [Commitments2023], with similar messages from as other sources [RSG2016, LSC+2023]. 

Meanwhile, regarding your question '*For example, if this classification model performs poorly on a certain category, can this interpretability guide me in adjusting the model structure to improve performance in those categories?*', there is a growing literature --- including our work --- on demonstrating interpretability leads to practical benefits across different settings. For example, [YCT+2024] and our work demonstrated that better internal interpretability leads to improved zero-shot segmentation, [MRM+2024] showed that one can leverage interpretability (based on sparse auto-encoders) to improve the accuracy of classifiers in language models, and [GES2024] demonstrated that one can use interpretability as a tool to remove spurious features.

Thank you again for your question on the motivation of building white-box models. Given the white-box/interpretable models are useful for building trustworthy and safe AI systems, our work aims to study whether it is possible to scale up such interpretable models and close their performance gap compared with black-box ones. Per your suggestion, we will add the above discussions on the practical benefits of interpretability to our revised version. 


[RSG2016] "Why Should I Trust You?": Explaining the Predictions of Any Classifier. Marco Tulio Ribeiro, Sameer Singh, Carlos Guestrin. KDD 2016.

[LSC+2023] On the importance of interpretable machine learning predictions to inform clinical decision making in oncology. Sheng-Chieh Lu, Christine L. Swisher, Caroline Chung, et al. Frontiers in Oncology, 2023.

[Commitments2023] Voluntary AI Commitments. https://www.whitehouse.gov/wp-content/uploads/2023/07/Ensuring-Safe-Secure-and-Trustworthy-AI.pdf 

[GL2010] Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. ICML 2010.

[CLW+2018] Theoretical Linear Convergence of Unfolded ISTA and its Practical Weights and Thresholds. Xiaohan Chen, Jialin Liu, Zhangyang Wang, Wotao Yin. NeurIP S2018.

[ZTA+2019] Deep Network Classification by Scattering and Homotopy Dictionary Learning. John Zarka, Louis Thiry, Tomás Angles, Stéphane Mallat. arXiv:1910.03561. 2019.

[YBP+2023] White-Box Transformers via Sparse Rate Reduction. Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin David Haeffele, Yi Ma. NeurIPS 2023.

[YCT+2024] Emergence of Segmentation with Minimalistic White-Box Transformers. Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, Yi Ma. CPAL 2024.

[GES2024] Interpreting CLIP's image representation via text-based decomposition. Yossi Gandelsman, Alexei A. Efros, Jacob Steinhardt. International Conference on Learning Representations (ICLR), 2024.

[MRM+2024] Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models. Samuel Marks, Can Rager, Eric J. Michaud, et al. https://arxiv.org/abs/2403.19647. 2024.
We are grateful for you engaging with our rebuttal further, and for your critical perspective on the work, which will no doubt improve it. Below we attempt to resolve the questions you posed.

> **Q10**: *At approximately 200G FLOPs, CRATE-α-L/16 achieves 84.6%, which means a further 7% increase is needed to match ViT's 85.2%. This gap should not be overlooked. The authors suggest that "this narrowing performance gap from Base to Large model size suggests that CRATE-α can nearly match ViT’s performance in large-scale settings"; however, it is observed in Figure 1 (Right) that as the model size increases, achieving such gains becomes increasingly challenging. For instance, CRATE-α-L/8 requires nearly double the FLOPs to achieve performance comparable to ViT-L/16. At smaller scales, CRATE-α appears to keep pace with ViT more easily, reaching comparable performance with a lower absolute increase in FLOPs. However, proportionally, CRATE-α consistently requires nearly double the increase in FLOPs to match ViT’s performance, as seen in the comparisons between CRATE-α-B/16 and ViT-B/32, as well as CRATE-α-L/14 and ViT-B/16.*

**A10**: We would like to clarify that **the performance gap between CRATE-α-L/16 (84.6% on IN1K) and ViT-L/16 (85.2% on IN1K) is 0.6%, instead of 7%**. We do agree that there is still a gap between CRATE-α-L/16 and ViT-L/16 (under similar FLOPs), but we would also like to highlight that this gap is much smaller than the previous gap between vanilla CRATE and ViT.

Meanwhile, to the best of our knowledge, there are no public results for ViT-L/8 on a similar setup. It is not clear whether increasing model size from ViT-L/16 to ViT-L/8 could significantly improve the performance of ViT. In particular, the FLOPs of the transformer architecture are closely related to the computation in the self-attention module, which grows quadratically with token length. When the patch size is halved, the token length increases to 4 times compared to the original one, resulting in the computational cost increasing by a factor of 16 for the self-attention module. We believe there is a trade-off between FLOPs consumption and accuracy gains in this context, and the transition from a patch size of 16 to 8 may not be optimal for CRATE-α and ViT.

If we first exclude the model with a patch size of 8 (due to no public available ViT-L/8), we can see that the gap on IN1K **narrows from 4.8%** between *CRATE-α-B/32 (76.5%) and ViT-B/32 (81.3%)* **to 0.6%** between *CRATE-α-L/16 (84.6%) and ViT-L/16 (85.2%)*. The latter group also has similar FLOPs and throughput. This is why we claimed that "this narrowing performance gap from Base to Large model size suggests that CRATE-α can nearly match ViT’s performance in large-scale settings."

We would like to highlight that the main focus of this paper is to compare the differences between CRATE-α and CRATE, and investigate whether it is possible to scale up white-box models and achieve competitive performance. We acknowledge that we did not conduct a comprehensive and thorough comparison between CRATE-α and ViT across various dimensions. Figure 1 (right) did not initially unify these variables, including patch size. Per your suggestion, we will add more experimental results (e.g., including different patch sizes for both ViTs and CRATE-α) in our revised version. While we think this is an interesting direction to explore in the future, it is beyond the main focus of this paper.
Thank you again for thoroughly reviewing our manuscript and response. We are grateful for your valuable feedback on our work, which will no doubt improve it.  We will add the results of supervised segmentation tasks to the main body of the revised version.
The reviewer appreciates the detailed response provided by the authors. Most of the concerns are addressed and the rating is upgraded accordingly. One suggestion that the authors could consider is adding the experiments of supervised training on segmentation tasks which is important to examine the dense prediction ability of white-box transformers.
Regarding A3 and A5, the authors claim that *"that under similar FLOPs, large CRATE-α-L/16 models nearly match ViT-L/16 performance."*  However, this assertion is not convincing when considered alongside Figure 1 (Right) and the response to 'Q1: Comparison with ViT' in the common response.

- At approximately 200G FLOPs, CRATE-α-L/16 achieves 84.6%, which means a further 7% increase is needed to match ViT's 85.2%. This gap should not be overlooked. The authors suggest that *"this narrowing performance gap from Base to Large model size suggests that CRATE-α can nearly match ViT’s performance in large-scale settings"*; however, it is observed in Figure 1 (Right) that as the model size increases, achieving such gains becomes increasingly challenging. For instance, CRATE-α-L/8 requires nearly **double** the FLOPs to achieve performance comparable to ViT-L/16.

- At smaller scales, CRATE-α appears to keep pace with ViT more easily, reaching comparable performance with a lower **absolute increase** in FLOPs. However, proportionally, CRATE-α consistently requires nearly **double** the increase in FLOPs to match ViT’s performance, as seen in the comparisons between CRATE-α-B/16 and ViT-B/32, as well as CRATE-α-L/14 and ViT-B/16.

Regarding A4, the authors reiterate that *"it is very possible to scale up white-box models that are designed from principles to achieve the same performance as black-box ones."* However, I would like to ask what practical **benefits** this white-box design (the so-called interpretability) has, aside from scaling up. For example, if this classification model performs poorly on a certain category, can this interpretability guide me in adjusting the model structure to improve performance in those categories? If the only advantage is scaling up, then such interpretability holds no real benefit, as black-box transformers can scale up more effectively (Figure 1 Right).

Regarding A6, the authors' explanation **conflicts** with the paper. Please refer to the last two lines of the Figure 1 caption: *"CRATE is trained only on ImageNet-1K, while CRATE-α and ViT are pre-trained on ImageNet-21K."* Why do the authors now explain that *"all four model variants are pretrained on IN21K and then fine-tuned on IN1K"*? I kindly request the authors to consider my question closely: since CRATE has not been pre-trained on ImageNet-21K while CRATE-α has been *"pre-trained
on ImageNet-21K and then fine-tuned on ImageNet-1K"* (as indicated in the caption of Table 1), how can these two models be compared fairly?
Thank you. I have no further comment at this point.