Scaling White-Box Transformers for Vision
Authors: Jinrui Yang
Xianhang Li
Druv Pai
Yuyin Zhou
Yi Ma
Yaodong Yu
Cihang Xie
Venue: NeurIPS 2024 poster (2024)
Keywords: white-box deep neural networks
representation learning
transformer
sparse coding
scaling
Primary area: deep_learning_architectures
TL;DR: We propose methods to scale white-box, transformer-like deep network architectures.

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.

Scaling White-Box Transformers for Vision
Jinrui Yang*l
Xianhang Li
Druv Pai
Yuyin Zhou' Yi Ma
Yaodong Yu Cihang Xie Fl
"qqual technique contribution, 'equal advising
UC Santa Cruz
UCC Berkeley
Abstract
CRATE, white-box transformer architecture designed to learn compressed and sparse representations, offers an intriguing alternative to standard vision trans- formers (ViT') due to its inherent mathematical interpretability. Despite extensive investigations into the scaling behaviors of language and vision transformers, the scalability of CRATE remains open question which this paper aims to address. Specifically, we propose CRATE-a, featuring strategic yet minimal modifications to the sparse coding block in the CRATE architecture design, and light training recipe designed to improve the scalability of CRATE. Through extensive experiments, we demonstrate that CRATE-a can effectively scale with larger model sizes and datasets. For example, our CRATE-Q-B substantially outperforms the prior best CRATE-B model accuracy ImageNet classification by 3.7%, achieving accuracy of 83.2%. Mcanwhile, when scaling further, our CRATE-Q-L obtains an ImageNet classification accuracy of 85.1%. More notably, these model performance im- provements are achieved while preserving. and potentially even enhancing the interpretability of learned CRATE models, we demonstrate through showing that the learned token representations of increasingly larger trained CRATE-a models yield increasingly higher-quality unsupervised object segmentation of images. The project page is https: htt///aayyjjaaagee werraaggggtttee
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, 5. 12]. However, the design of the Transformer architecture and its many variants remains largely empirical and lacks rigorous mathematical interpretation. This has largely hindered the development of new Transformer variants with improved efficiency interpretability. The recent white-box Transformer model CRATE [46] addresses this gap by deriving simplified Transformer block via unrolled optimization 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, I1)) of the features. As shown in previous work [47], besides mathematical interpretabil- ity, the learned CRATE models and features also have much better semantic interpretability than conventional transformers, i.c., visualizing features 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 pathway to improved performance and emergent proper- ties [44, 40, 41, 14]. Until now, the deployment of CRATE has been limited to relatively modest scales. The most extensive model described date is the base model size encompasses 77.6M parameters
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

85
1216
CRATE-B/32
68.5
Top- Accuracy on ImnggNNt-1K
D/16
L/14
US
Overcomplete Dictionary
+5.3
80
1932
Decoupled
Dictionary
+2.0
75
* CRATE-U
CRATE
Residual
+0.7
76,5
L/16
VIT
R/16
60
70
80
70
100
300
400
Top-1 Accuracy ImageNet-1K
FLOPs
Figure 1: (Left) We demonstrate how modifications to the components enhance the performance of the CRATE model. The four models are trained using the same setup: first pre-trained on ImageNet-21K and then fine-tuned on ImageNet- K. Details are provided in Section 3. (Right). We compare the FLOPs and accuracy on ImageNet-IK of our methods with ViT [9] and CRATE [46]. The values of CRATE-o: model correspond to those presented in Table 1. A more detailed comparison between CRATE-a and VIT included Appendix A.2.
(CRATE-Large) [46]. This contrasts sharply with standard Vision Transformers (VITs [9]), which have been effectively scaled 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, Le., Tiny, Small, Base. Large, Huge. Detailed model specifications 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
limiting factor that hinders further scaling. To overcome this, significantly expand the channels, decouple the association matrix, and add a residual connection, resulting in a new model variant CRATE-O. 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-a architecture.
provide extensive experiments supporting the effective scaling of our CRATE-a models. For example, we scale the CRATE-Q model from Base to Large size for supervised image classification on ImageNet-211 21K [6], achieving 85. % top-1 accuracy on ImageNet- I at the Large model size. We further scale the model size from Large to Huge, utilizing languae pre-training with contrastive learning on DataComp1 [10], and achieve a zero-shot top-/ accuracy of 72.3% on ImageNer-IK the Huge model size. These results demonstrate the strong scalability of the CRATE-O model, shedding light on scaling up mathematically interpretable models for future work.
The main contributions of this paper are threefold:
We design three strategic yet minimal modifications for the CRATE model architecture unleash its potential. In Figure 1, we reproduce the results of the CRATE model within our training setup, initially pre-training on ImageNet-211
classification and subsequently fine-tuning on ImageNet- IK classification. Compared to the vanilla CRATE model that achieves 68.5% top-1 classification accuracy on ImageNet-IKK our CRATE-Q-B/27 BB32 model significantly improves the vanilla CRATE model by 8%, which clearly demonstrates the benefits of the three modifications to the existing CRATE model. Moreover, following the settings of the best CRATE model and changing the image patch size from 32 to 8, our CRATE-a-B model attains top-1 accuracy of 83.2% ImageNet- IK exceeding the previous best CRATE model's 's score 79.5% by significant margin 3.7%.
2. Through extensive experiments, show that one can effectively scale CRATE-a via model size and data simultaneously. model size, there is marginal improvement top-1 classification accuracy (+0.5%, from 70.8% to 71.3%) on ImageNet- indicating diminished returns [46]. Furthermore, by scaling the training dataset, achieved substantial 1.9% improvement in top-1 classification accuracy ImageNet- increasing from 83.2% to 85.1 % when going from CRATE-A Base to Large.
contrast, when increasing the CRATE model from Base to Large
3. We further successfully scale CRATE-a model from Large to Huge leveraging vision-languee language pre-training on DataComp1 Compared to the Large model, the Huge model (CR (CRATE-a-H) achieves zero-shot top-1 classification accuracy of 72.3% ImageNet-1K, marking significant
Model configurations detailed Table (in Appendix A.1).

---

scaling gain of 2.5% over the Large model. These results indicate that the CRATE architecture has the potential serve as an effective backbone for vision-language foundation models.
Related Work
White-box Transformers. [46, 45] argued that the quality of a learned representation can be assessed through unified objective function called the sparse rafe reduction. Based on this framework, [46, 45] developed a family of transformer-likk 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, 45J. 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 VIT's in terms of model size and data size. From the perspective of self- supervision, MAE [12] provides a scalable approach to effectively training a ViT-Huge model using only ImageNet-IK: IK. 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-1 achieving consistent performance improvements with the scaling of ViT model size. From the perspective of supervised learning. [49, 5] present comprchensive analysis of the empirical scaling laws for vision transformers on image classification tasks, sharing some similar conclusions with [15]. [49] suggests that the performance- compute 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 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.
2 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 [21..... x E RR DxN to represent the data, where each æ; E R is "token", such that each data point is a realization of X. For instance, X can represent a collection of image patches for image, and æ; the i-th image patch. We use f E F: D
RdxM denote [21,.... ... zN| E d denote the E IRd denotes the feature of the i-th input token a, The
the mapping induced by the transformer, and we let Z = f(x)
features for input data X. Specifically, transformer consists of multiple, say L, layers, and so can be written f f= f ..... f frre, where f N + d denotes the E-th layer of the transformer, and the pre- pre-processing layer denoted by free 1 DxN RdxN The input the f-th layer f of the transformer is denoted by z = . y E R so that fe: Z + zo+1 In particular, Z' = (X) E Rd N denotes the output of the processccnng layer and the input to the first layer.
2.1 Sparse Rate Reduction
Following the framework proposed in [45], we posit that the goal of representation learning is to learn a feature mapping representation f F: RDxN + _x 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 union of low-dimensional subspaces, say with orthonormal bases U = (U+)ke(K) E (RdxP) K_ [46] proposes the Sparse Rate Reduction (SRR) objective to measure the goodness of such a learned representation:
max Ez= (Lar(())) = min Ez=)(x) Je
- R(Z UlK)) + +
(I)
fet
3

---

where Z f(X) denotes the token representation. Z denotes the l' norm, and R(Z),
( UpR) are (estimates for) rate distortions [4, 7]. defined as:
R(Z) = log det (I+
1
d 2 No 2)
IZ ((K) = Σ R(U 2).
km
In particular, RZ U(K1) (resp. R(Z) provide closed-form estimates for the number of bits required to encode the sample up to precision e > 0, conditioned (resp. unconditioned) on the samples being drawn from the subspaces with bases Ujx]. Minimizing the term R 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 representation learning via the rate reduction approach.
CRATE: Coding RATE Transformer
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 a proximal gradient descent step for optimizing learning objective L(Z), i.e., Z t+1 f'(Z') Prox[Z
V (())) Here we use 17 to denote the step size and Prox[-] denote the proximal operator [27].
One layer 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 f-th layer of CRATE is defined as
Z f" (z*) ISTA(Z (ZTTAZZ+1/22 11/2 D³), where 2 Z' MSSA Z²).
Compression block (MSSA). The compression block in CRATE, called Multi- Subspace Self- Attention block (MSSA), derived for compressing the token set Z [z1..... zx] by optimizing the compression term R2 (defined Eq. (1)), i.e..
2-1122 Z MS U UkK) Z
zV U(k));
where U denotes the (local) signal model layer e, and the MSSA operator is defined
[UT Z) (allmaxx([]] Z) (UT Z))
kp N
MSSA( Z UKK]) =
UU ... Uk]
(5)
(U Z) softmax (U 2) (U Z))
Compared with the commonly used attention block in transformer [42], where the k-th attention head is defined (V, Z) softmax((Q) (Q] Z) (K, 2)), MSSA uses only matrix obtain the query, key, and value matrices in the attention: that is, U. = Q = K₁-- = Vg.
Sparse coding block (ISTA). The Iterative Shrinkage- Thresholding Algorithm (ISTA) block is designed to optimize the sparsity term and the global coding rate term, Al| R(Z Ujk) in (1), [46] shows that an optimization strategy for these terms posits a (complete) incoherent dictionary D ddd and takes a proximal gradient descent step towards solving the associated LASSO
problem arg min | 2 2-1/2
Dd Z A|| [[]], obtaining the iteration
zz IZ D ReLU Z +1/2 (D*) (Z D² n)1). (6)
In particular, the ISTA block sparsifies the intermediate iterates Z 2+1/2 w.r.t. D to obtain z
3 CRATE-a Model
In this section, we present the CRATE-a architecture, which a variant of CRATE [46]. As shown in Fig. 1 (Right), there is significant performance gap between the white-box transformer CRATE-B/16

---

z
A
ISTA-overcomplete
D
A'
ODL
z'
A Proe, Af (D') z 22)
Z 1+1/2
ODL
MSSA
Atteationileadl:.
MSSA
21
Aggregate
Attectionieadll...
Z
Figure 2: One layer of the CRATE-O 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 provided in Fig. the Appendix.
(70.8%) and the vision transformer ViT-B/16 (84.0%) [9]. One possible reason is that the ISTA block applics a complete dictionary D E Idxdd which may limit its expressiveness. In contrast, the MLP block in the transformer' applies two linear transformations W1, W2 Rdx4d leading to the MLP block having times more parameters than the ISTA block.
Since the ISTA block in CRATE applies 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 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 complete dictionary D E IR dxd we use an overcomplete dictionary D' E pdx(Cd), where C > I (a positive integer) is the overcompleteness parameter. Furthermore, we explore two additional modifications 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 atten- tion block Z 1211/2 we propose to sparsify the token representations with respect to an overcomplete dictionary Dd (Cd) by optimizing the following LASSO problem,
A RJ arg min
Z 2
1+/2 - A 2 + A|A|Il]
A20
To approximately solve (7), we apply two steps of proximal gradient descent, i.e.,
A = O. A = Prox,, a [A; D' 2
A Prox,, A D', 2 2
(8)
where Prox the proximal operator the above non-negative LASSO problem (7) and defined as Prox,, [A; D, Z ReLUUA- DDD (DA Z) - n)1).
The output of the sparse coding block is defined ze+1
A'A where = ===, ISTA-0C D
(10)
Namely, A is sparse representation of 2+11/2 with respect to the overcomplete dictionary D The original CRATE ISTA tries to learn complete dictionary D E R to transform and sparsify the features Z. By leveraging more atoms than the ambient dimension, the overcomplete dictionary D (Cd) can provide redundant yet expressive codebook to identify the salient sparse
"The MIP block defined z z + W20 z(+1/2), where is the nonlinear activation function and f+1/2 denotes the output of the attention block.
5

---

Training Loss ImageNet-214 (Varying Model Sizes)
Training Loss ImageNet-21k (Varying Patch Sizes]
7.0
7.0
CRATE-a-L/32
6.5
6.5
CRATE-a-L/14
Training Loss
6.0
Training Loss
6.0
CRATE-a-
5.5
5.5
5.0
5.0
CRATE-a-T/16
CRATE-a-S/16
4.0
4.0
CRATE-a-B/16
3.5
CRATE-a-L/14
3.5
3.0
3.0
50000
100000
150000
200000
250000
50000
100000
150000
200000
250000
Step
Step
Figure 3: Training loss curves of CRATE-a on ImageNet-21K, (Leeft) Comparing training loss curves across CRATE-Q with different model sizes. (Right) Comparing training loss curves across CRATE-Q- Large with different patch sizes.
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 De in the last step (defined in (10) of the sparse coding block, z D' AA where D' Rdx(Cd) is a different dictionary compared to D By introducing the decoupled dictionary, we further improve the model performance by 2.0%, shown in Fig. 1. We denote this mapping from 22+1/2 to ze+1 as the Overcomplete Dictionary Learning block (ODL), defined as follows:
ODL(Z 6-1/2 D D') IST 1122
D'==
(11)
Modification #3: Residual connection. Based on the previous two modifications, we further add residual connection, obtaining the following modified sparse coding block:
z = 2d ODL 11+ D',
(12)
An intuitive interpretation of this modified sparse coding block as follows: instead of directly sparsifying the feature representations Z, we first identify the potential sparse patterns present in Z by encoding it over learned dictionary. Subsequently, we incrementally refine Z by exploiting the sparse codes obtained from the previous encoding step. From Fig. find that the residual connection leads to a 0. 0.7% improvement.
To summarize, effectively scale white-box transformers, implement three modifications the vanilla white-box CRATE model proposed in [46]. Specifically, in our CRATE-a model, we introduce decoupling mechanism, quadruple the dimension of the dictionary (4x), and incorporate 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-o from Base to Large size and ultimately to Huge size. (2) Downstream applications: To further verify the broader benefits of scaling the CRATE-O model, we conduct additional experiments on real-world downstream tasks and present preliminary exploration results of CRATE-Q language tasks. (3) Interpretability: In addition to scalability, study the interpretability of CRATE-Q across different model sizes.
4.1 Dataset and Evaluation
Scaling Study. For the transition from Base to Large size, we pre-train our model on ImageNet-21KK and fine-tune on ImageNet- IK via supervised learning. When scaling from Large Huge, we utilize the DataCompl [10] dataset within a vision-language pre-training paradigm, allowing us to study the effects of scaling the model to massive size. For evaluation, we evaluate the zero-shot accuracy of these models on ImageNet-IK.

---

Table I: Top-1 accuracy of CRATE-Q on ImageNet-1K with different model scales when pre-trained on ImageNet-21K K and then tune on ImageNet- IK. For comparison, we also list the results from the paper [46] which demonstrate the diminished return from CRATE base large, trained only ImageNet- KK. "IN-21K" refers to ImageNet-21K. (*Results from [46].)
Models (Base)
ImageNet- (K6%)
Models (Large)
ImageNet- (K(%)
CRATE-B/16 w/ IN-21K
70.8*
CRATE-L/16 w/o IN-21K
71.3
CRATE-a-B/32
76.5
CRATE-a-LAA 0/332
80.2
CRATE-a-B/16
81.2
CRATR-O-I/HL
83.9
CRATE-a-B/8
83.2
CRATE-a-L/8
85.1
Downstream Applications. We include additional experimental results on four downstream datasets CCFFAR-10100 Oxford Flowers, and Oxford-IIT Pets). We also examine the dense prediction capa- bility of CRATE-a by training it on segmentation tasks using the ADE20K dataset [51]. For language tasks, we conduct new experiments with CRATE-a using autoregressive training on Open WebText, 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.
4.2 Training and Fine-tuning Procedures
Scaling Study. (1) Base to Large size: We initially pre-train the CRATE-a model on ImageNet-21 and subsequently fine-tune it on ImageNet- IK. During the pre-training phase, we set the learning rate to 10 4, 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 1.6 x 10-4, maintain weight decay at 0.1, and batch size at 4096. We apply label smoothing with 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 224x224. In both the pre-training and fine-tuning phases, we use the AdamW optimizer [24] and incorporate warm-up strategy, characterized by linear increase over 10 epochs. Both the pre-training and fine-tuning are conducted for 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 84x84, 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 the image size. A smaller image size results in a faster training speed. In the configurations of BRAATE---CLIP-R CRATE-Ox 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.c., the CRATE-Q model. As we will show in the later results, this setup effectively demonstrates the scaling behaviors of CRATE-a models in the image domain. Detailed hyperparameter settings can be found in Appendix A.
CRATE-O-CLIPA-L. and PRATE-O-CLIPA-I we use the
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-Q 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 following the setup in [16]. We compare CRATE-Q models (57M and 120M) with CRATE and GPT-2, using results from CRATE reported in [45].
CRATE-a using autoregressive training on WebnWebTeee
7

---

CLIPA Training Loss on DataComp1B
Zero-shot Accuracy on ImageNet-22 1k
CRATE-a-B/16
80
Pre-train
79.1
Training Loss
3.5
CRATE-a-H/14
CRATE-a-L/14
Top-1 Accuracy
70
Fine-tune
69.8
66.7
72.2
63.6
60.8
80
2.0
52.9
50-
40
0.5
10000
20000
30000 40000
50000
60000
70000
30
GR1 CRATE-G-UI4 CRATE-
VIT-H/14
Step
Model
Figure 4: (Left) Comparing training loss curves of CRATE-Q-CLIPA with different model sizes on DataComp1B: IB. (Righr) Comparing zero-shot accuracy of CRATE-a-B/L/H models and VYT-H on ImageNet-1K
4.3 Results and Analysis
Scaling the CRATE-Q Model from Base to Large. As shown in Table we compare CRATE-Q-B and CRATE-Q-L at patch sizes 32, 16, and 8. Firstly. we find our proposed CRATE-Q-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 B to CRATE-L results in only 0.5% improvement on ImageNet-1K. This indicates a case of diminishing returns. These findings compellingly highlight that the scalability of CRATE-G models significantly outperforms that of the vanilla CRATE. Meanwhile, the training loss the pre-training stage 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-Q Model from Large to Huge. From the results shown in Fig. 4, we find that: (I) MERATE---CLIPA-LAA significantly outperforms CRAAEE--CLIPPA-BIIII by 11.3% and 9.0% in terms of ImageNet- IK zero-shot accuracy during the pre-training and fine-tuning stages, respectively. The substantial benefit suggests that the quality of learned representation may 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-Q- BRATE-O-CLIP-HIE CLIP-HH114 continues to benefit from larger training datasets, outperforming CRATE-a-CLIP-LL14 by 3.1% and 2.5% in terms of ImageNet- K zero-shot accuracy during the pre-training and fine-tuning stages, respectively. This demonstrates the strong scalability of the CRATE-a model. To explore the performance ceiling, we train a standard ViT-CLIPA-H/I4 from scratch and observe improved performance.
Downstream Applications. On four downstream datasets, as shown in Table 2, we find that CRATE-a consistently outperforms CRATE, with both models pre-trained on IN21K, while CRATE-a demonstrates improved performance as model size increases. For the segmentation task, results in Table show that CRATE-a consistently outperforms CRATE across key metrics, with both models pre- pre-trained on I.21K. These findings indicate significant performance gains in vision tasks the language task, Table shows that CRATE-a significantly improves over CRATE in language modeling. Due to limited time and resource constraints, we completed 80% of the total iterations for CRATE-a-small and 55% for CRATE-u-base. compared the 600K total iterations used for CRATE. Nevertheless, CRATE-a still demonstrated notable improvements.
beyond classification.
Interpretability. As shown in Fig. we provide the segmentation visualization on COCO val2017 [20] for CRATE-Q, CRATE, and ViT, respectively. We find that our model preserves and even improves the (semantic) interpretabilityy advantages of CRATE. Moreover, we summarize quan- titative evaluation results on COCO val2017 in Table 6. Interestingly, when scaling up model size for CRATE-a, the Large model improves over the Base model terms of object detection and segmentation.
4.4 Compute-efficient Scaling Strategy
We further explore methods to scale models efficiently in terms of computation. Table demonstrates that the CRATE-a model scales effectively from the Base model its larger variants. However, the pre-training computation for the top-performing model, CRATE-a-L/8, is resource-intensive

---

Table 2: The performance comparison between CRATE and CRATE-a various datasets.
Dataset
CRATE-B/32
CRATE-a-B/32
CRATE-a-L/32
CRATE-a-B/16
CRATE-a-L/14
CIFAR-I0
97.22
98.17 17
98.68
98.67
99.10
CIFAR-100
97.77 85.27
91.16 89.40
90.58
92.57
Oxford Plowers-102
99.01
99.27
99,56
Oxford-III -
80.38
88.19
90.46
92.70
93.98
Table 3: Performance comparison of CRATE models with different configurations.
Model
Scope
mIoU
mAcc
aAcc
CRATE-a-B/32
global 35.35
45.28
77.63
CRATE-B/32
global 30.28
39.29
75.21
Table 4: The comparison between CRATE and CRATE-a on the NLP task using the Open Web dataset.
GPT-2-base CRATE-base
CRATE-a-small
CRATE-a-basc
Model size
124M
60M
57M
120M 3.14
Cross-entropy validation loss
2.85
3.37
3.28
Table 5: Compute-efficien scaling strategy. To reduce the compute requirements of the pre-trainn pre-training stage, we use model with a larger patch size. This results a shorter token length for the same input size. The second and fourth columns indicate the compute requirements for the pre-training and fine-tuning stages, respectively, measured in TPU v3 core-hours. Details are provided in Section 4.4.
Pre-train
Core-hours
Fine-tune
Core-hours
Total core-hours
IN-IK((%)
CRATE-a-L/14
872
3,524
83.7
CRATE-Q-L/32
2,652
CRATE-Q-L/8
3,486
6,138
84.2
CRATE-Q-L/14
8,947
CRATE-Q-L/14
872
9,819
83.9
CRATE-a-L/8 35,511
CRATE-a-L/8
3,486
38,997
8.5.1
ImageNet-2IK. 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 summarized in Table 5.
Results and analysis. (1) When fine-tuning with CRATE-Q-L/14 and using CRATE-a-L/32 for pre- training on ImageNet-21K, this approach consumes about 35% of the TPU v3 core-hours required by CRATR-O-L/14, yet achieves a promising 83.7% top-1 accuracy on ImageNet- K, comparable the 83.9% achieved by CR CRATE-O-L/14; (2) When fine-tuning with CRATE-a-L/8 and using CRATE-O-L//2 for pre-training, this approach consumes just 15% of the training time required by CRATE-a-L/8. yet it still achieves promising 84.2% top-1 accuracy on ImageNet-IK, compared to 85.1% when using the CRATE-a-L/8 model the pre-training stage; (3) While the total computational cost of CRATE-a-L/32 + CRATE-a-L/8 is less than that of CRATE-a-L/14 CRATE-a-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-Q models in the future.
5 Discussion
Limitations. Although have used some existing compute-cfficient training methods (c.g.. CLIPA [18]) and have initiated an exploration into compute-cfficient scaling strategies for white- box transformers in Section 4.4, this work still requires 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 growing interest in developing white-box transformers for better interpretability and transparency across wide range of tasks and domains, including image segmentation [46], self-supervised masked autoencoders [26), and integrated sensing and communications [50], etc. Moreover, results the scalability of white-

---

CRATE-U
CRATE
VIT
Figure 5: Visualization of segmentation COCO val2017 [20] with MaskCut [43]. (Top row)
Supervised CRATE-a effectively identifies the main objects in the image. Compared with CRATE (Middle row), CRATE-a achieves better segmentation performance in terms of boundary. (Bottom row) Supervised ViT fails to identify the main objects in most images. We mark failed image with
Table 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-O model achieves notable performance gain. In addition, when scaling CRATE-a from base large, also exhibits the benefit of scalability.
Detection AP7s
Segmentation AP7s f
Model
Train
AP50 f
AP f
AP50
AP
CR CRATE-B/8 [47] VIT-B/8 [47]
Supervised Supervised
2.9
1.0
1.3
2.2
0.7
1.0
0.8
0.2
0.4
0.7
0.5
0.4
CRATE-a-BAA
Supervised
3.5
1.1
1.5
2.2
1.0
1.1
CRATE-a-L/8
Supervised
4.0
1.7
2.0
2.7
1.1
1.4
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 wide range of applications and tasks, benefiting broad audience interested in building more interpretable and performant deep learning models and further amortizing the pre-training compute costs.
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-a model along both 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-lis like architectures.
Acknowledgement
This work is supported by a gift from Open Philanthropy, TPU Research Cloud (TRC) program, and Google Cloud Research Credits program.
10

---

References
[1] Thomas Blumensath and Mike Davies. Iterative thresholding for sparse approximations. Journal of Fourier analysis and Applications, 14:629-655 66629-654, 2008.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877--9011, 2020.
[3] Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Theoretical linear convergence of unfolded ista and its practical weights and thresholds. Advances in Neural Information Processing Systems, 31, 2018.
[4] Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
[5] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pages
PMLR, 2023.
[6] Jia Deng, Wei Dong. Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
[7] Harm Derksen, Yi Ma, Wei Hong, and John Wright. Segmentation of multivariate mixed data via lossy coding and compression. In Visual Communications and Image Processing 2007, volume 6508, pages 170-181. SPIE, 2007.
[8] Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1810.04805, 2018
[9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 20X1v:2010.11929, 2020.
[10] Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smymis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, et al. Datacomp: In search of the next generation of multimodal datasets. Advances in Neural Information Processing Systems, 36, 2024.
[11] Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th international conference on international conference on machine learning, pages 399 406. 2010.
[12] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVE conference on computer vision and putlern recognition, pages 16000-16009, 2022.
[13] Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong. John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openclip. July 2021.
[14] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint a.X0v:2310.06825, 2023.
[15] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, arXiv preprint orXiv:2001.0834 2020.
[16] Andrej Karpathy. nanogpt, 2022. URL https: htt.ttb..cm//arpppttt/nnnof GitHub repository.

---

[17] Mingyang Li, Pengyuan Zhai. Shengbang Tong, Xingjian Gao, Shao-Lun Huang, Zhihui Zhu, Chong You, Yi Ma, al. Revisiting sparse convolutional model for visual recognition. Advances in Neural Information Processing Systems, 35:10492-10500 -0504, 2022.
[18] Xianhang Li, Zeyu Wang, and Cihang Xie. An inverse scaling law for clip training. Advances in Neural Information Processing Systems, 36, 2024.
[19] Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language-image pre-training via masking. In CVPR, 2023.
[20] Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740-755. Springer, 2014.
[21] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[22] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11976- 1986, 2022.
[23] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In ICLR, 2017.
[24] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint a7X1v:1711.05101,,00 2017.
[25] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2018.
[26] Druv Pai, Ziyang Wu, Sam Buchanan, Tianzhe Chu, Yaodong Yu, and Yi Ma. Masked completion via structured diffusion with white-box transformers. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
[27] Neal Parikh, Stephen Boyd, et al. Proximal algorithms. Foundations and trends in Optimiza- fion, 1(3):127-239, 2014.
[28] Qing Qu, Yuexiang Zhai, Xiao Li, Yuqian Zhang, and Zhihui Zhu. Geometric analysis of nonconvex optimization landscapes for overcomplete learning. In International Conference on Learning Representations, 2019.
[29] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-8888 8763. PMLR, 2021.
[30] Sucheng Ren, Fangyun Wei, Zheng Zhang, and Han Hu. Tinymim: An empirical study of distilling mim pre- tre-trained models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3687-3697, 2023.
[31] Mannat Singh. Quentin Duval, Kalyan Vasudev Alwala, Haoqi Fan, Vaibhav Aggarwal, Aaron Adcock, Armand Joulin, Piotr Dollár, Christoph Feichtenhofer, Ross Girshick, al. The effec- tiveness of mae pre- prepretraining for billion-scale pretraining. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 5484 5494, 2023.
[32] Quan Sun, Yuxin Fang, Ledell Wu, Xinlong Wang. and Yue Cao. Eva-clip: Improved training techniques for clip at scale. arXiv preprint arXiv:2303. 00003./5300 15389, 2023.
[33] Quan Sun, Jinsheng Wang, Qiying Yu, Yufeng Cui, Fan Zhang, Xiaosong Zhang, and Xinlong Wang. Eva-clip-1880 18b: Scaling clip to 18 billion parameters. arXiv preprint arXiv:2402.0422 2024.
12

---

[34] Xiaoxia Sun, Nasser M Nasrabadi, and Trac D Tran. Supervised deep sparse coding networks. In 2018 25th IEEE International Conference on Image Processing (ICIP), pages 346-350, IEEE, 2018.
[35] Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1-9, 2015.
[36] Bahareh Tolooshams and Demba Ba. Stable and interpretable unrolled dictionary learning. arXiv preprint arXiv:2/06.00058, 2021.
[37] Ily Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mip-mixer: An all-mlp architecture for vision. Advances neural information processing systems, 34: 24261-24272, 2021.
[38] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347-10357. PMLR, 2021.
[39] Hugo Touvron, Matthieu Cord, and Hervé Jégou. Deit ii:: Revenge of the vit. In European conference on computer vision, pages 516-533. Springer, 2022.
[40] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1399 2023.
[41] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.092 2023.
[42] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[43] Xudong Wang, Rohit Girdhar, Stella Yu, and Ishan Misra. Cut and learn for unsupervised object detection and instance segmentation. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 3124-3134, 2023.
[44] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022.
[45] Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Hao Bai, Yuexiang Zhai, Benjamin D Haeffele, and Yi Ma. White-box transformers via sparse rate reduction: Compression all there is? arXiv preprint 1rXiv:2311.131100 2023.
[46] Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin Haeffele, and Ma. White-box transformers via sparse rate reduction. Advances in Neural Information Processing Systems, 36, 2023.
[47] Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, and Yi Ma. Emergence of segmentation with minimalistic white-box transformers. In Conference on Parsimony and Learning, pages 72-93. PMLR, 2024.
[48] John Zarka, Louis Thiry, Tomás Angles, and Stéphane Mallat. Deep network classification by scattering and homotopy dictionary learning. arXiv preprint arXiv: 190.03191. 2019.
[49] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- ers. Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pages 2210-11211,, 2022.
13

---

[50] Bowen Zhang and Geoffrey Ye Li. White-box arXiv:2407.02251, 2024.
for isac. arXiv preprint
[51] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 633-641, 2017.
14

---

Appendix
A Additional Experiments and Details
A.1 Model configuration.
provide details about CRATE-o model configurations in Table
Table 7: Model configurations for different sizes of CRATE-a, parameter counts, and comparisons to CRATE models. is depth, d the hidden size, and K is the number of heads.
Model Size
K
CRATE-Q # Params
CRATE # Params
Tiny
12
192
3
4.8M
1.7M
Small
12
576
12
41.0M
13.1M
Base
12
768
12
72.3M
22.8M
Large
24
1024
16
253.8M
77.6M
Huge
32
1280
16
526.8M
159.8M
Table 8: The comparison between CRATE-a and ViT. FLOPs and throughput are calculated based on an input size of 224x224 on NVIDIA RTX A6000 graphics card.
Model
FLOPs
#Params
Throughput
Model
FLOPs
#Parums (M)
Throughput
CRATE-a-B/32
6.4
74.0
499
VIT-B/32
4.4
88.2
706
CRATE-a-B/16
25.8
72.3
233
VIT-B/16
17.6
86.5
375
CRATE-a-L/32
22.8
256.0
215
VIT-L/32
15.4
306.5
329
CRATE-a-1/14
119.7
253.7
VIT-/14
304.1
A.2
Comparison of model structure with VIT.
We also compare CRATE-a VIT in terms of computational costs, number of parameters, and inference speed. These comparisons are summarized in Table 8, where CRATE-G matches ViT's effi- ciency while achieving similar accuracy. With the same number of layers and embedding dimensions, CRATE-a has fewer parameters than VIT, and its FLOPs/Througpuut is slightly higher.
To more accurately compare CRATE-a and VIT with larger model sizes, we conduct experiments on CRATE-a-L/16 with an image resolution of 336, nearly matching the setup of ViT-L/16. Both models similar amount of FLOPs: 210G for CRATE-Q-L/I L/16 compared to 191G for ViT-L/16. The throughput, or images processed per second, is also comparable at 35.53 for our model versus 35,56 for ViT-I//16. The accuracy of CRATF-O-L//6 reach 84.6%, closely approaching VIT's 85.2% under similar conditions. Meanwhile, combining the trend from Figure/ (right) in the main paper, this narrowing performance gap from Base to Large model size suggests that CRATE-G nearly matche ViT's performance in large-scale settings. Besides, CRATE-a inherits the mathematical interpretability of the white-box models and can also achieve much better semantic interpretability evaluated by zero-shot segmentation.
A.3 Training details of CRATE-Q-CLIPA models.
When employing the CRATE-a architecture to replace the vision encoder in the CLIPA [18] framework. we essentially follow the original CLIPA training recipe. The setup for the pre-training stage is presented in Table 9. During the fine-tuning stage,
made some modifications: the input image size set to 224 x 224, the warmup steps are set to 800, and the base learning rate set to 4e-7. When calculating the loss, the last token from the text encoder as the text feature.
the classification token from the vision encoder the image feature and
To explore the performance ceiling, also train ViT-CLIPA model from scratch. Most of the hyperparameters remain the same as those in Table but there are some modifications the pre- training stage. The batch size is set to 65,536, and the text length is set to 8 to speed up training. As
15

---

with the CLIPA setup, wp steps are set to 3,200. Additionally, we add color jitter and grayscale augmentation, and use global average pooling instead of the classification token. These modifications help stabilize training.
Config
Value
optimizer optimizer momentum
AdamW [25]
(0.9, 0.95)
batch size base Ir minimal warm-up steps schedule
32768
8c-6
0
1600
cosine decay [23]
weight decay random crop area
0.2
(40, 100)
resize method
bi-linear
temperature init
1/0.07 [13, 19]
Table 9: Pre-training hyper-parameetes for CLIPA.
Modification decoupled dictionary
Modification #1 ovexcomplete
g
z
Tatk-wwereeee lata
D
A'
overccooppe.... Loto
ODL
ODL
Kn -ene X-1 vizi 10 monj
/-1/2
ODL
z
ODL
Modification
residual connection
z
Figure 6: One layer of the CRATE-a model architecture (with more details for the three modifications described in Section 3.
16

---

Visualization of self-attention maps of CRATE-G. We provide visualization of attention maps of CRATE-Q Fig. 7.
Figure 7: We visualize the self-attention maps of the CRATE-a Base model using 8 x 8 patches trained using classification. Similar the original CRATE [47], our model also demonstrates the capability to automatically capture the structural information of objects. For each row, the original image is displayed on the left, while the corresponding self-attention maps are shown on the right. The number of self-attention maps corresponds to the number of heads in the CRATE-a model.
Visualization of loss curves. We visualize the training loss curves of the four models, including CRATE and its three variants, in Fig. 8. We visualize the training loss curves of CRATE-a-Band with different patch sizes in Fig. 9. In Fig. 10, we also visualize the training loss curves of models trained with efficient scaling strategy described in Section 4.4 in the main paper.
Training Loss on ImageNet- 21k (Varying Model Sizes)
8.0
CRATE
7.5
CRATE (+0)
Training Loss
7.0
CRATE (+0+D)
CRATE (+0+D+R)
6.5
CRATE-Q
6.0
5.5
5.0
4.5
50000
100000
150000 Step
200000
250000
Figure Training loss curves of different model architectures (mentioned Fig. in the main paper) on ImageNet-21KK. The patch size is 32 for all four models shown in this figure. (+O: +overcomplete dictionary, +D: +decoupled dictionary, +R: +residual connection.)
17

---

Training Loss on ImageNet- 21k (Varying Patch Sizes)
7.0
CRATE at-
6.5
CRATE CRATE-a-B/16
Training Loss
6.0
BRATE-aa-Bp
5.5
5.0
4.5
4.0
3.5
50000
100000
150000
200000
250000
Step
Figure 9: Comparing training loss curves across CRATE-Q-Base Base with different patch sizes.
Training Loss (with Efficient Scaling)
4.50
4.25
Training Loss
4.00
3.75
3.50
3.25
CRATE-O-L/32 (pre-train)
3.00
CRATE-o-L/32->11 14 (fine-tune)
2.75
CRATE-a-L/32->8 ffine-tune)
2,50
200000
220000
240000
260000
280000
300000
Step
Figure 10: Comparing training loss curves when using the efficient scaling strategy. The blue curve corresponds to the CRATE-O-Large/32 model (in the pre-training stage). After pre-training the CRA/E-o-Lagee32, we further fine-tune with smaller patch sizes (longer token length), including patch size 14 (orange curve) and patch (green curve).
18

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA]. .
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
general, answering
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly outline the primary contributions and scope of the paper. For specific details, refer to abstract and Sections 1, which align with the claims stated.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
19

---

Justification:
We discuss the limitations of the work in section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approxinations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide details of reproducing the experimental results in Section 4 and the appendix.
Guidelines:
20

---

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 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 specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
httbbmmissiooooicyyy...... for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //n ceaaps..c/public/uuii hetpoldeggggid///eeeeeeee /CodeSubmisse ismmisssionn ssssnnolicyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
21

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: tative.
We evaluate our method on a large public validation dataset, which is represen-
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, 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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful 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 the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
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.
22

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics attps://neures
Answer: [Yes]
Justification: we follow the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: we discuss them in Section 5.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [NA]
Justification: Our work has potential risk for misuse.
23

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All assets used in the paper are properly credited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperswi eesorsswitcodeee thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: we haven't 't introduced new assets in this work.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?

---

Answer:
[NA]
Justification:
This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
25


Reviews:
[{"review_id": "BkVqYtV1dW", "reviewer": "Reviewer_8F4z", "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.", "questions": "see weakness", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a careful study to enhance the performance of CRATE. The paper introduces key modifications to the existing CRATE, including the sparse coding block, decoupled dictionary, and residual connection. \n- The paper investigates its scaling behavior and shows promising improvements of the newly introduced CRATE-alpha.\n- The paper presents in-depth experiments, such as the scaling analysis on ImageNet. The paper also shows improvements for semantic interpretability. \n- The figures and model architecture are well-illustrated.", "weaknesses": "Overall I find the paper is well-presented and solid. Below are my minor concerns for this paper:\n- 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.\n- It would be interesting to further understand its potential downstream applications (not only vision but also language data)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below we attempt to resolve the questions you posed.\n\n>**Q1**: *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.*\n\n**A1**: We agree that some of our findings mainly apply to white-box models (e.g., white-box transformers like CRATE [YBP+2023]). 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 [YBP+2023], self-supervised masked autoencoders [PBW+2024], and integrated sensing and communications [ZL2024], etc. We believe that our findings and insights could be helpful for developing white-box transformers for a wide range of applications and tasks. Moreover, our results on the scalability of white-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 [GL2010, SNT2018, CLW+2018, ZTA+2019, DLZ+2022] designed via unrolled optimization. In summary, we believe that the findings and insights of this work could benefit a broad audience in the NeurIPS community interested in building more interpretable and performant deep learning models.\n\nThank you for your suggestion on discussing the impact of this work, and we add the above discussion to our revision.\n\n\n>**Q2**: *It would be interesting to further understand its potential downstream applications (not only vision but also language data)*\n\n**A2**: Thank you for your suggestion. Please refer to our response to '**Q2: Additional experimental results on real-world downstream tasks**'  and '**Q3: Performance of CRATE-α on NLP task**' in the common response.\n\nWe again thank you for your review, and hope we have provided satisfactory responses to your questions. Please let us know if you have further questions or comments.\n\n[YBP+2023] White-Box Transformers via Sparse Rate Reduction. Yaodong Yu, Sam Buchanan, Druv Pai, et al. NeurIPS 2023.\n\n[PBW+2024] Masked Completion via Structured Diffusion with White-Box Transformers. Druv Pai, Sam Buchanan, Ziyang Wu, et al. ICLR 2024.\n\n[ZL2024] White-Box 3D-OMP-Transformer for ISAC. Bowen Zhang, Geoffrey Ye Li. arXiv:2407.02251.\n\n[GL2010] Learning fast approximations of sparse coding. Karol Gregor and Yann LeCun. ICML 2010.\n\n[SNT2018] Supervised deep sparse coding networks. Xiaoxia Sun, Nasser M Nasrabadi, Trac D Tran. ICIP 2018.\n\n[CLW+2018] Theoretical Linear Convergence of Unfolded ISTA and its Practical Weights and Thresholds. Xiaohan Chen, Jialin Liu, Zhangyang Wang, et al. NeurIPS 2018.\n\n[ZTA+2019] Deep Network Classification by Scattering and Homotopy Dictionary Learning. John Zarka, Louis Thiry, Tomás Angles, et al. arXiv:1910.03561. 2019.\n\n[DLZ+2022] Revisiting sparse convolutional model for visual recognition. Xili Dai, Mingyang Li, Pengyuan Zhai, et al. NeurIPS 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nLdISuoKMG", "reviewer": "Reviewer_yemz", "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.", "questions": "Apart from the questions in weakness above, another question is:\n\nwhy the performance in dense prediction tasks is so bad?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. It is quite meaningful to study white-box transformers and try to increase their scalability which promises its application in potential usage.\n\n2. Comprehensive evaluation. The proposed method is validated on multiple datasets and tasks which demonstrate the scalability of CRATE-$\\alpha$.\n\n3. The presentation is clear. Overall, the paper is well-organized and the method is easy to follow.", "weaknesses": "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.\n\n2. 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?\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below we attempt to resolve the questions you posed.\n\n>**Q1**: *Performance gaps with vanilla ViT. As shown in Figure 1, CRATE-α 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.*\n\n**A1**: Thank you for your suggestions on comparison between ViT and CRATE-α. Please refer to our response to '**Q1: Comparison with ViT**' in the common response. \n\n>**Q2**: *According to the model configuration, the number of parameters of CRATE-α 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-α? Is it simply because of larger models?*\n\n**A2**: Considering the number of parameters, CRATE-α-B/16 (72.3M) is comparable to CRATE-L/16 (77.6M), as shown in table 4 of the Appendix. CRATE-α-B/16 still surpasses CRATE-L/16 by 9.9% (81.2% vs 71.3%) in ImageNet-1K (see Table 1 of the manuscript). Meanwhile, from Table 1 of the manuscript, we also observe that increasing the size of CRATE from base (22.8M) to large (77.6M) results in only a 0.5% improvement, demonstrating diminished returns at the cost of 3.4 times more parameters. In contrast, we found that change from CRATE-B/16 to CRATE-α-B/16 leads to a much more significant improvement. Moreover, we also tried directly scaling CRATE-L to a larger dataset (ImageNet-21K), but the experiment showed it is challenging to train CRATE-L directly due to the instability of the optimization, and its performance is even worse than just using ImageNet-1K. \n\nTherefore, combining these observations, we can find that the design of CRATE-α indeed leads to improved performance than CRATE, which is not simply due to the number of parameters being larger. Without a carefully designed architecture, naively increasing parameters might not lead to improved performance. Thank you for your question, and we will add the above discussion to our revised version.\n\n>**Q3**: *Although the authors made lots of efforts in scaling CRATE to CRATE-α, 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.*\n\n**A3**: Thank you for your suggestion on the interpretability part of the paper presentation. *The main focus and motivation* of this paper is to investigate whether it is possible to scale up white-box deep learning models to achieve competitive performance as black-box ones such as ViTs, which is a *long-standing problem for white-box deep learning models* [GL2010, ZTA+2019, DLZ+2022, YBP+2023]. Our work *provides an affirmative answer* to this question. \n\nMeanwhile, the previous work [YBP+2023, PBW+2024] has already explored the mathematical interpretability of white-box transformers. Our newly designed white-box model CRATE-α follows the white-box design principles as previous work, and the operators and architectures are derived by optimizing the sparse coding rate reduction Eq. (1) in the manuscript. Therefore, the mathematically interpretability [YBP+2023, PBW+2024] developed in previous work is inherited by this work. It would be highly interesting to provide a more comprehensive study on the interpretability of  CRATE-α, such exploration is out-of-scope for the present work and left for future work.\n\nPer your suggestions, we have conducted more experiments on the zero-shot segmentation (following the same setup in Section 4.4), including segmentation visualization on more images as shown in Figure 5. We will add these new experimental results to the appendix of our revision.\n\n>**Q4**: *Why the performance in dense prediction tasks is so bad?*\n\n**A4**: The reported results in Table 3 of the manuscript are evaluated **under the zero-shot segmentation setting**. For the results in Table 3, we were using MaskCut [WGY+2023], a self-supervised segmentation method, to study the interpretability of different models (ViT, CRATE, CRATE-α). In particular, we directly extracted features using the learned models and applied MaskCut to perform segmentation. We did not perform supervised training on segmentation tasks. Therefore, the segmentation results are not as competitive as the standard ones which are typically trained with segmentation labels. Thank you for your suggestion. We will add a paragraph and highlight this difference in our revised version.\n\nWe again thank you for your review, and hope we have provided satisfactory responses to your questions. Please let us know if you have further questions or comments.\n\n[WGY+2023] Cut and learn for unsupervised object detection and instance segmentation. Xudong Wang, Rohit Girdhar, Stella X Yu, et al. CVPR 2023.\n\n[GL2010] Learning fast approximations of sparse coding. Karol Gregor and Yann LeCun. ICML 2010.\n\n[ZTA+2019] Deep Network Classification by Scattering and Homotopy Dictionary Learning. John Zarka, Louis Thiry, Tomás Angles, et al. arXiv:1910.03561. 2019.\n\n[DLZ+2022] Revisiting sparse convolutional model for visual recognition. Xili Dai, Mingyang Li, Pengyuan Zhai, et al. NeurIPS 2022.\n\n[YBP+2023] White-Box Transformers via Sparse Rate Reduction. Yaodong Yu, Sam Buchanan, Druv Pai, et al. NeurIPS 2023.\n\n[PBW+2024] Masked Completion via Structured Diffusion with White-Box Transformers. Druv Pai, Sam Buchanan, Ziyang Wu, et al. ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jxtodcLWcI", "reviewer": "Reviewer_cPqJ", "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.", "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.\n   \n2. 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?\n   \n3. 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?\n   \n4. 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*.\n   \n5. How was the amount of training data determined? Was there a specific standard or a FLOPs constraint? For example:\n   - 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.\n   - In Section 4.1, for training models from Large to Huge, there were **2.56** billion and **512** million training samples, respectively.", "limitations": "The authors adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "**Originality:** The paper continues the white-box design philosophy of the original CRATE model while integrating advanced techniques such as overparameterized sparse coding, decoupled dictionary, and residual connections. Although some of these techniques have been previously validated, successfully combining them with a white-box Transformer is a noteworthy achievement. The integration not only works effectively but also yields commendable results.\n\n**Quality:** The paper is technically sound overall, employing rigorous notation and formula definitions to elucidate the design principles. The proposed models demonstrate significant improvements compared to the previous generation of CRATE models. Additionally, the authors are careful and honest in evaluating the weaknesses and limitations of their work.", "weaknesses": "**Clarity:**\n- 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.\n- The organization of Section 4 is not concise, making it easy for readers to lose track.\n  - 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.\n  - The frequent interleaving of experimental setup descriptions with the presentation of experimental results disrupts the flow and coherence of the narrative.\n\n**Significance:** \n- 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. \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below we attempt to resolve the questions you posed.\n\n>**Q1**: *The paper is heavily symbolized, ... it severely hampers understanding of the paper's details ...*\n\n**A1**: Thank you for the paper presentation suggestion. We've added a new diagram to our rebuttal pdf (Figure 1 in the rebuttal pdf) illustrating the three modifications. In the revised version, we'll include additional figures/diagrams to clarify these modifications and enhance explanations in Section 3, improving our presentation.\n\n>**Q2**: *The organization of Section 4 is not concise, making it easy for readers to lose track.* \n\n**A2**: Thank you for your suggestion on the paper presentation. In our revision, we will reorganize Section 4 to further improve the clarity, including providing an overview of the datasets used in the paper in a paragraph, and adding pointers to different figures/tables when introducing the datasets.\n\n>**Q3**: *Although CRATE-α shows significant improvements over the original CRATE model, it still lags behind the state-of-the-art.*\n\n**A3**: Please see our response to '**Q1: Comparison with ViT**' in the common response, indicating that under similar FLOPs, large CRATE-α-L/16 models nearly match ViT-L/16 performance.\n\n>**Q4**: *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.*\n\n**A4**: We agree that performance/accuracy is an important metric for evaluating different models. However, we believe that interpretability is also an important property in machine learning models built for real-world applications, where reliability and trustworthiness are paramount, such as self-driving and robotics (among a long list of others). As we have shown in this work (including new experimental results described in **A3**), it is very possible to scale up white-box models that are designed from principles to achieve the same performance as black-box ones. Since, as we have shown, we can train white-box models with significantly improved mathematical and semantic interpretability properties compared to the usual transformer at negligible performance drop, we believe that our model may be well-suited for real-world applications in which such properties are necessary.\nWe will add a paragraph on discussing the trade-off between interpretability and performance to Section 5 Discussion in our revised version. \n\n> **Q5**: *As previously mentioned, as shown on the right side of Figure 1, CRATE-α usually requires approximately twice the FLOPs to reach the performance level of ViT, putting it at a noticeable disadvantage.*\n\n**A5**: see our response in **A3**.\n\n>**Q6**: *How does the performance improvement of CRATE-α 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?*\n\n**A6**: In Figure 1 (left), **all four model variants are pretrained on IN21K and then fine-tuned on IN1K**. We can see that CRATE-α-B/32 (76.5% on IN1K) outperforms CRATE-B/32 (68.5% on IN1K) by 8.5%. We also tried training the original CRATE-L model on 21K, but the training was unstable and we found it is challenging to scale CRATE-L to IN21K. The CRATE-L model cannot be effectively scaled to larger datasets due to unstable optimization. This is also the motivation for proposing the CRATE-α model.\n\n>**Q7**: *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?*\n\n**A7**: Thank you for your suggestion. We have plotted the validation loss for Figure 3, the validation loss curves follow the similar trend as the training loss curves in Figure 3. We will add these new results to our revised version.\n\n>**Q8**: *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-α-L/32 + CRATE-α-L/8 versus CRATE-α-L/14 + CRATE-α-L/14.*\n\n**A8**: Thank you for your suggestions regarding the comparison. Given the higher number (10x) of training samples in pre-training compared to fine-tuning, pre-training dominates compute costs. Following your advice, we've summarized the total cost of pre-training and fine-tuning in Table 4 of our rebuttal pdf file. 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. This reinforces the benefit of reducing image token sequence lengths during pre-training for compute-efficient scaling.\n\n> **Q9**: *How was the amount of training data determined? Was there a specific standard or a FLOPs constraint?*\n\n**A9**: Regarding the number of training epochs, we mainly follow the setup in previous work: [YCT+2024], [YBP+2023], and [TCJ2022]. Specifically, C.1 section in paper [YCT+2024], Table 13 of C.5 in paper [YBP+2023], and section 4.1 in paper [TCJ2022]. For the number of training samples from large to huge, we referred to Table 1 of [LWX2023].\n\nWe again thank you for your review, and hope we have provided satisfactory responses to your questions. Please let us know if you have further questions or comments.\n\n[YCT+2024] Emergence of Segmentation with Minimalistic White-Box Transformers. Yaodong Yu, Tianzhe Chu, Shengbang Tong, et al. CPAL 2024.\n\n[YBP+2023] White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is? Yaodong Yu, Sam Buchanan, Druv Pai, et al. arXiv:2311.13110. 2023.\n\n[TCJ2022] DeiT III: Revenge of the ViT. Hugo Touvron, Matthieu Cord, Hervé Jégou. ECCV 2022.\n\n[LWX2023] An Inverse Scaling Law for CLIP Training. Xianhang Li, Zeyu Wang, Cihang Xie. NeurIPS 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BpLwbrZzBm", "reviewer": "Reviewer_YH3U", "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.", "questions": "See the weakness.", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a novel architecture, CRATE-α, that builds upon the existing CRATE model with minimal yet strategic modifications, enhancing scalability without compromising interpretability.\n\nThe authors provide a wealth of empirical evidence supporting the effectiveness of CRATE-α, including comparative results on ImageNet classification tasks and a detailed analysis of training behaviors across different model scales.\n\nA key strength is the paper's focus on maintaining the interpretability of the model, which is often a trade-off in scaling deep learning models. The authors demonstrate that CRATE-α models retain high-quality unsupervised object segmentation capabilities.\n\nThe paper includes a thorough exploration of scaling behaviors, from Base to Large to Huge model sizes, using both supervised learning on ImageNet and vision-language pre-training with contrastive learning on DataComp1B.", "weaknesses": "Could the proposed architecture work well on other tasks like NLP?\n\n\nWhile 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below we attempt to resolve the questions you posed.\n \n>**Q1**: *Could the proposed architecture work well on other tasks like NLP?*\n\n**A1**: Thank you for your suggestions on new experiments on NLP. Please refer to our response to '**Q3: Performance of CRATE-α on NLP task**' in our common response.\n\n>**Q2**: *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.*\n\n**A2**: Thank you for your suggestions on new experiments on other datasets and real-word applications. Please refer to our response to '**Q2: Additional experimental results on real-world downstream tasks**' in our common response.\n\nWe again thank you for your review, and hope we have provided satisfactory responses to your questions. Please let us know if you have further questions or comments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
### **Common response to all reviewers**:

We thank all reviewers for their insightful feedback. We are especially encouraged by their recognition of:

- The novelty and impact of our central ideas (` Reviewer YH3U `: “The paper presents a novel architecture, CRATE-α, …, enhancing scalability without compromising interpretability.” ` Reviewer cPqJ `: “successfully combining them with a white-box Transformer is a noteworthy achievement. The integration not only works effectively but also yields commendable results.”)
- The benefits of the new white-box architecture we have proposed (` Reviewer 8F4z `: “shows promising improvements of the newly introduced CRATE-alpha.”; ` Reviewer yemz `: “Comprehensive evaluation.”; ` Reviewer cPqJ `: “The proposed models demonstrate significant improvements compared to the previous generation of CRATE models.”; ` Reviewer YH3U `: “The authors provide a wealth of empirical evidence supporting the effectiveness of CRATE-α”)
- The quality of the exposition (` Reviewer 8F4z `: “The figures and model architecture are well-illustrated.”; ` Reviewer yemz `: “the paper is well-organized and the method is easy to follow.”)
- The motivation of our work (` Reviewer yemz `: “It is quite meaningful to study white-box transformers and try to increase their scalability …”)

In the remainder of this message, we address certain questions raised by the reviewers.

> ### **Q1: Comparison with ViT**

**A1**: We agree that the proposed white-box CRATE-α is not significantly outperforming ViT on classification tasks. However, based on the newly designed dictionary learning block and training recipe, we were able to significantly improve over the vanilla white-box CRATE and nearly closed the gap between white-box transformers and ViTs, especially when the model size becomes larger. 

**Performance comparison.** In Figure 1 (right), ViT was fine-tuned at 384x384 resolution, while CRATE was at 224x224. To more accurately compare CRATE-α and ViT with larger model sizes, we conducted experiments on CRATE-α-L/16 with an image resolution of 336, nearly matching the setup of ViT-L/16. Both models used a similar amount of FLOPs: 210G for CRATE-α-L/16 compared to 191G for ViT-L/16. The throughput, or images processed per second, was also comparable at 35.53 for our model versus 35.56 for ViT-L/16. The accuracy of CRATE-α-L/16 reached 84.6%, closely approaching ViT’s 85.2% under similar conditions. Meanwhile, combining the trend from Figure 1 (right), this narrowing performance gap from *Base* to *Large* model size suggests that CRATE-α can nearly matche ViT’s performance in large-scale settings. Besides, CRATE-α inherits the mathematical interpretability of the white-box models and can also achieve much better semantic interpretability evaluated by zero-shot segmentation.

**Efficiency comparison.** We would like to thank ` Reviewer yemz ` for suggestions on comparison with ViT in computational costs, number of parameters, and inference speed. These comparisons are summarized in the Table 1 of our rebuttal pdf, where CRATE-α matches the efficiency of ViT while achieving similar accuracy. With the same layers/embedding dimension, CRATE-α has fewer parameters than ViT, and the FLOPs/Throughput of CRATE-α is slightly higher than ViT. 

We will add the above discussion and new experimental results to our revised version.

> ### **Q2: Additional experimental results on real-world downstream tasks**

**A2:** We have conducted new experiments on supervised segmentation tasks and fine-tuning on additional downstream datasets.

**Segmentation.** We compare the performance of CRATE and CRATE-α for the segmentation task on ADE20K dataset to study the benefits of the newly designed architecture CRATE-α compared to CRATE. Following the setup of [RWZ+2023], we compared them in segmentation tasks, focusing on direct comparisons due to time and resource constraints, without extensive parameter tuning. Results demonstrate 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. We will include these results in our revised version.

| Model        | Scope  | mIoU  | mAcc  | aAcc  |
| ------------ | ------ | ----- | ----- | ----- |
| CRATE-α-B/32 | global | 35.35 | 45.28 | 77.63 |
| CRATE-B/32   | global | 30.28 | 39.29 | 75.21 |

**Fine-tuning classification datasets.** We included additional experimental results on evaluations on downstream datasets for CRATE-α. From the table 2 of our rebuttal pdf, we find that CRATE-α consistently outperforms CRATE (both models are pre-trained on IN21K), and CRATE-α achieves improved performance when model size increases.

We will add the above new results to our revised version.

> ### **Q3: Performance of CRATE-α on NLP task**

**A3**: Thanks to ` Reviewers 8F4z ` and ` Reviewer YH3U ` for suggesting we explore CRATE-α in language tasks. We conducted new experiments for CRATE-α with autoregressive training on OpenWebText, following the setup in [Karpathy2022]. We compare CRATE-α models (57M and 120M) with CRATE and GPT-2, where the result of CRATE is from [YBP+2023]. Table 3 in our rebuttal pdf shows that CRATE-α significantly improves over CRATE in language modeling. Due to limited time and resource constraints, compared to the total number of iterations (600K) used in CRATE, we only finished 80% of the total iterations for CRATE-α-small and 55% for CRATE-α-base. CRATE-α still showed notable improvements. We will add these new results to our revised version.

[Karpathy2022] NanoGPT. https://github.com/karpathy/nanoGPT. Andrej Karpathy

[YBP+2023] White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is? Yaodong Yu, Sam Buchanan, Druv Pai, et al. arXiv:2311.13110. 2023.

[RWZ+2023] TinyMIM: An Empirical Study of Distilling MIM Pre-Trained Models. Sucheng Ren, Fangyun Wei, Zheng Zhang, et al. CVPR 2023.

Meta-review: 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.
Decision: Accept (poster)