SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left\| \left(\boldsymbol {\Delta} - \alpha \boldsymbol {Z} _ {t} \left(\boldsymbol {I} + \alpha \boldsymbol {Z} _ {t} ^ {*} \boldsymbol {Z} _ {t}\right) ^ {- 1} \left(\boldsymbol {Z} _ {t} ^ {*} \boldsymbol {\Delta} + \boldsymbol {\Delta} ^ {*} \boldsymbol {Z} _ {t}\right)\right) \left(\boldsymbol {I} + \alpha \boldsymbol {Z} _ {t} ^ {*} \boldsymbol {Z} _ {t}\right) ^ {- 1} \right\| _ {\mathrm {F}} \leq \frac {9}{4} \| \boldsymbol {\Delta} \| _ {\mathrm {F}}, \tag {114}
$$
which gives the claim after taking suprema.
B Additional Experiments and Details
In this section, we provide details about our experiments, and report the results of additional experiments that were not covered in the main text. CRATE takes arguably the most basic design choices possible, and so we do not attempt to directly compete with state-of-the-art performance from heavily engineered and empirically designed transformers. The results of our experiments are meant to convey a few core messages:
- Despite not being engineered to compete with the state-of-the-art, CRATE performs strongly on large-scale real-world datasets, including classification on ImageNet-1K. CRATE also achieves strong transfer learning performance.
- Because our model is designed through unrolled optimization of a well-understood objective, each layer is interpretable. In particular, we can analyze the performance of CRATE, as well as design network modifications, on a layer-wise basis. This is powered by an arguably unparalleled level of insight into the role of each operator in our network.
- We make the simplest possible choices during the design of CRATE, but these can be changed easily while keeping the same framework. We study a few modifications later in this section (Appendix B.4) and show that they do not significantly hurt empirical performance, but emphasize here that there is significant potential for improvement with different architecture choices (and in particular a different theoretical analysis).
In this subsection, we provide more details for implementing CRATE on vision tasks.
B.1.1 Architecture of CRATE
Architectural modifications. Compared to the conceptual architecture proposed in Sections 2.5 and 3, we make the following change for the sake of implementation simplicity:
- In the compression step, replace the term  $\frac{p}{N\epsilon^2} [U_1,\dots ,U_K]$  in the MSSA operator with another trainable parameter  $\mathbf{W}\in \mathbb{R}^{d\times pK}$ . Thus the MSSA block becomes
$$
\operatorname {M S S A} \left(\boldsymbol {Z} \mid \boldsymbol {U} _ {[ K ]}, \boldsymbol {W}\right) \doteq \boldsymbol {W} \left[ \begin{array}{c} \operatorname {S S A} (\boldsymbol {Z} \mid \boldsymbol {U} _ {1}) \\ \vdots \\ \operatorname {S S A} (\boldsymbol {Z} \mid \boldsymbol {U} _ {K}) \end{array} \right]. \tag {115}
$$
PyTorch code for CRATE. We provide PyTorch-style code for implementing our proposed network architecture. Algorithm 1 defines the overall architecture, Algorithm 2 and Algorithm 3 contain details for the transformer block, self-attention block (MSSA-block), and MLP block (ISTA-block).
B.1.2 Training Setup
Pre-training on ImageNet-1K. We apply the Lion optimizer [73] for pre-training both CRATE and ViT models. We configure the learning rate as  $2.4 \times 10^{-4}$ , weight decay as 0.5, and batch size as 2,048. We incorporate a warm-up strategy with a linear increase over 5 epochs, followed by training the models for a total of 150 epochs with cosine decay. For data augmentation, we only apply the standard techniques, random cropping and random horizontal flipping, on the ImageNet-1K dataset. We apply label smoothing with smoothing parameter 0.1. One training epoch of CRATE-Base takes around 240 seconds using 16 A100 40GB GPUs.
Fine-tuning. We fine-tune our pre-trained CRATE and ViT models on the following target datasets: CIFAR10/CIFAR100 [10], Oxford Flowers-102 [7], Oxford-IIIT-Pets [16]. We also evaluate our pre-trained models on the commonly used ImageNet Real [36] benchmark. For each fine-tuning task, we use the AdamW optimizer [26]. We configure the learning rate as  $5 \times 10^{-5}$ , weight decay as 0.01, and batch size to be 512. To allow transfer learning, we first resize our input data to 224. For data augmentations, we also adopt several standard techniques: random cropping, random horizontal flipping, and random augmentation (with number of transformations  $n = 2$  and magnitude of transformations  $m = 14$ ).<sup>14</sup>
$^{14}$ https://github.com/huggingface/pytorch-image-models/blob/main/timm/data/auto AUGMENT.py
Algorithm 1: PyTorch-style pseudocode for CRATENetwork
Class ViT dictionary definition   
CRATE: #initialization def init(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool  $=$  'cls', channels  $= 3$  ,dim_head  $= 64$  ,dropout  $= 0$  .emb_dropout  $= 0$  ). #define patch, image dimensions and number of patches image_height, image_width  $\equiv$  pair(image_size) patch_height, patch_width  $\equiv$  pair(patch_size) numpatches  $\equiv$  (image_height // patch_height) \* (image_width // patch_width) patch_dim  $\equiv$  channels \* patch_height \* patch_width #define patch embedding, positional embedding, dropout, and transformer self.to_batch_embedding  $\equiv$  Sequential(Rearrange, LayerNorm(patch_dim), Linear(patch_dim, dim), LayerNorm(dim)) self(pos_embedding  $\equiv$  Parameter(random(1, num_batches + 1, dim)) self.xls_token  $\equiv$  Parameter(random(1, 1, dim)) self.dropout  $\equiv$  Dropout(emb_dropout) self.transformer  $\equiv$  Transformer(dim, depth, heads, dim_head, mlp_dim, dropout) #define pooling, latent layer, and MLP head self.pool  $\equiv$  pool self.tolatent  $\equiv$  Identity() self.mlp_head  $\equiv$  Sequential(LayerNorm(dim), Linear(dim, num_classes)) #forward pass def forward(self, img): x = self.to_batch_embedding(img) b,n,_  $=$  shape(x) clf_tokens  $\equiv$  repeat(self.cls_token,'1 1 d->b 1 d',b = b) x  $=$  concatenate((clf_tokens,x),dim=1) x += self(pos_embedding[:,:(n + 1)] x = self.dropout(x) x = self.transformer(x) x = mean(x, dim = 1) if self.pool == 'mean' else x[:,0] x = self.to_forent(x) return self.mlp_head(x)
Algorithm 2: Pytorch Style Pseudocode for Transformer Block in CRATE
# Class Transformer definition   
class Transformer: #initialization def init(self，dim，depth，heads，dim_head，mlp_dim，dropout  $= 0$  ： #define layers self.layers  $=$  [] self.depth  $=$  depth for_in range(depth): self.layers.append([LayerNorm(dim，Attention(dim，heads，dim_head, dropout))]) self.layers.append([LayerNorm(dim，FeedForward(dim，mlp_dim， dropout))]) #forward pass def forward(self，x): for attn，ff in self.layers:  $\mathrm{x}_{-} = \mathrm{attn}(\mathrm{x}) + \mathrm{x}$ $\mathbf{x} = \mathbf{f}\mathbf{f}(\mathbf{x}_{-})$  return x
Algorithm 3: Pseudocode for Attention and FeedForward
Class FeedForward definition   
class FeedForward: #initialization def init(self, dim, hidden_dim, dropout = 0., step_size=0.1, lambda=0.1): self.weight = Parameter(Tensor(dim, dim)) init.kaiming.uniform_(self.weight) self.step_size = step_size self.lambd = lambd #forward pass def forward(self, x): x1 = linear(x, self.weight, bias=None) grad_1 = linear(x1, self.weight.t(), bias=None) grad_2 = linear(x, self.weight.t(), bias=None) grad_update = self_STEP_size * (grad_2 - grad_1) - self_STEP_size * self.lambd output = relu(x + grad_update) return output # Class Attention definition   
class Attention: #initialization def init(self, dim, heads = 8, dim_head = 64, dropout = 0.): inner_dim = dim_head * heads project_out = not (heads == 1 and dim_head == dim) self_heads = heads self_scale = dim_head ** -0.5 self attends = Softmax(dim = -1) self_dropout = Dropoutdropout) self.qkv = Linear(dim, inner_dim, bias=False) self.to_out = Sequential(Linear(inner_dim, dim), Dropout Dropout)) if project_out else nn.Identity() #forward pass def forward(self, x): w = rearrange(self.qkv(x), 'b n (h d) -> b h n d', h = self_heads) dots = matmul(w, w.transpose(-1, -2)) * self_scale attn = self.attn(dots) attn = self_dropout(attn) out = matmul(attn, w) out = rearrange(out, 'b h n d -> b n (h d)') return self.to_out(out)
B.2 Experimental Results
In this subsection, we provide additional experimental results on CRATE, including layer-wise measurements, visualizations, as well as ablation studies.
B.2.1 Layer-wise Evaluation and Visualization
Layer-wise evaluation of compression and sparsity. Similar to Figure 3, we conduct the layerwise evaluation of compression term and sparsity for CRATE-Tiny, CRATE-Base, and CRATE-Large. We observe similar behavior as mentioned in Section 3.1: both the compression term and the sparsity term improves as the layer index increases.
(a) Compression (Model: CRATE-Tiny).
(b) Sparsity (Model: CRATE-Tiny).
(c) Compression (Model: CRATE-Base).
(d) Sparsity (Model: CRATE-Base).
(e) Compression (Model: CRATE-Large).
(f) Sparsity (Model: CRATE-Large).
Figure 5: Left: The compression term  $R^{c}(Z^{\ell +1 / 2})$  of the MSSA outputs at different layers. Right: the sparsity of the ISTA output block,  $\| Z^{\ell +1}\| _0 / (d\cdot N)$ , at different layers.
Visualizing layer-wise token representations. In Figure 6, we visualize the token representations  $Z^{\ell}$  at different layers  $\ell \in \{1,\dots ,12\}$ . We provide more results evaluated on other samples in Appendix B.2.2.
Visualizing layer-wise subspaces in multi-head self-attention. We provide the visualization of  $U_{[K]}^{\ell}$  in Figure 7.
(a)  $\ell = 1$
(b)  $\ell = 2$
(c)  $\ell = 3$
(d)  $\ell = 4$
(e)  $\ell = 5$
(f)  $\ell = 6$
$(\mathbf{g})\ell = 7$
(h)  $\ell = 8$
(i)  $\ell = 9$
(j)  $\ell = 10$
$(\mathbf{k})\ell = 11$
Figure 6: Visualizing layer-wise token  $Z^{\ell}$  representations at each layer  $\ell$ . To enhance the visual clarity, we randomly extract a  $50 \times 50$  sub-matrix from  $Z^{\ell}$  for display purposes. (Model: CRATE-Tiny)
(1)  $\ell = 12$
$(\mathbf{k})$ $\ell = 11$
Figure 7: We visualize the  $[U_1^\ell, \ldots, U_K^\ell]^*[U_1^\ell, \ldots, U_K^\ell] \in \mathbb{R}^{pK \times pK}$  at different layers. The  $(i,j)$ -th block in each sub-figure corresponds to  $(U_i^\ell)^* U_j^\ell$  for  $i,j \in [K]$  at a particular layer  $\ell$ . To enhance the visual clarity, for each subspace  $U_i$ , we randomly pick 4 directions for display purposes. (Model: CRATE-Tiny)
B.2.2 Additional Layer-wise Visualization
We provide more results of the layer-wise token representation visualization on other samples in Figure 8, Figure 9, Figure 10, and Figure 11 (Model: CRATE-Base).
Figure 8: Visualizing layer-wise token  $Z^{\ell}$  representations at each layer  $\ell$ . To enhance the visual clarity, we randomly extract a  $50 \times 50$  sub-matrix from  $Z^{\ell}$  for display purposes. (Sample 1)
Figure 9: Visualizing layer-wise token  $Z^{\ell}$  representations at each layer  $\ell$ . To enhance the visual clarity, we randomly extract a  $50 \times 50$  sub-matrix from  $Z^{\ell}$  for display purposes. (Sample 2)
Figure 10: Visualizing layer-wise token  $Z^{\ell}$  representations at each layer  $\ell$ . To enhance the visual clarity, we randomly extract a  $50 \times 50$  sub-matrix from  $Z^{\ell}$  for display purposes. (Sample 3)
Figure 11: Visualizing layer-wise token  $Z^{\ell}$  representations at each layer  $\ell$ . To enhance the visual clarity, we randomly extract a  $50 \times 50$  sub-matrix from  $Z^{\ell}$  for display purposes. (Sample 4)
B.3 CRATE Ablation
Hyperparameters of CRATE. In Table 2, we present evaluation of CRATE trained with various parameters. More specifically, we investigate the effect of number of epochs, weight decay, learning rate, step size  $(\eta)$  and the regularization term  $(\lambda)$  in ISTA block. As shown in Table 2, CRATE demonstrates consistently satisfactory performance across a diverse range of hyperparameters.
Table 2: Top 1 accuracy of CRATE on various datasets with different architecture design variants when trained on ImageNet.
Model	epoch	weight decay	lr	η (ISTA)	λ (ISTA)	ImageNet
CRATE-B	150 (default)	0.5 (default)	2.4 × 10-4	0.1	0.1	70.8
CRATE-B	150	0.5	2.4 × 10-4	0.02	0.1	70.7
CRATE-B	150	0.5	2.4 × 10-4	0.5	0.1	66.7
CRATE-B	150	0.5	2.4 × 10-4	0.1	0.02	70.8
CRATE-B	150	0.5	2.4 × 10-4	0.1	0.5	70.5
CRATE-B	90	0.5	2.4 × 10-4	0.1	0.1	69.5
CRATE-B	300	0.5	2.4 × 10-4	0.1	0.1	70.9
CRATE-B	150	1.0	2.4 × 10-4	0.1	0.1	70.3
CRATE-B	150	0.05	2.4 × 10-4	0.1	0.1	70.2
CRATE-B	150	0.5	4.8 × 10-4	0.1	0.1	70.2
CRATE-B	150	0.5	1.2 × 10-4	0.1	0.1	70.3
B.4 Exploring Architecture Variants
In this section, we explore the two following alternative architectures. One architecture involves a modification to the attention mechanism, while the other involves a modification to the sparsification mechanism. Again, we re-emphasize that these choices, although principled, are entirely modular and the choices we make here still lead to very simple architectures. A more sophisticated analysis may lead to different, more complicated architectures that perform better in practice. The architectures we experiment with are:
- Compression-inspired attention mechanism: revert the change in (115). That is, the attention mechanism implements (11) and (12) directly.
- Majorization-minimization proximal step sparsification: instead of (17), implement (92).
We obtain the following classification results in Table 3. After conducting additional simplifications to the network architecture (i.e., imposing additional constraints to the network architecture design), we discover that CRATE maintains reasonable performance on ImageNet-1K.
Table 3: Top 1 accuracy of CRATE on various datasets with different architecture design variants when trained on ImageNet.
Model	MSSA-block	ISTA-block	ImageNet
CRATE-B	default	default	70.8
CRATE-B	Eq. (11) and (12)	default	63.3
CRATE-B	default	Eq. (92)	68.6
B.5 Sparse Coding vs. Non-Negative Sparse Coding
In the main body, we used a non-negative sparse coding formulation (16) to obtain the ISTA block as an unrolled proximal gradient step in (17). Suppose that we hadn't done this, and instead directly computed an ISTA block using an unrolled proximal gradient step on (15). Such a block would give the following update rule:
$$
\boldsymbol {Z} ^ {\ell + 1} = S _ {\lambda \eta} \left(\boldsymbol {Z} ^ {\ell + 1 / 2} + \eta \boldsymbol {D} ^ {*} \left(\boldsymbol {Z} ^ {\ell + 1 / 2} - \boldsymbol {D} \boldsymbol {Z} ^ {\ell + 1 / 2}\right)\right), \tag {116}
$$
where  $S_{\lambda \eta}$  is the soft-thresholding function
$$
S _ {\lambda \eta} (x) = \operatorname {s g n} (x) \cdot (| x | - \lambda \eta) _ {+}. \tag {117}
$$
applied element-wise to its input matrix. The resulting architecture would be an alternative to CRATE; below, we discuss some empirical and theoretical similarities and differences between the two formulations and architectures.
Empirical evaluation of soft-thresholding-based architecture. We summarize the results of CRATE with soft-thresholding activation in Table 4. We use  $\lambda = 10$  in  $S_{\lambda \eta}$  and set all other hyperparameters the same as in the original CRATE-Base evaluation on ImageNet-1K. We find that such a soft-thresholding model achieves slightly worse performance-a drop of  $3.2\%$  top-1 accuracy-compared to the default CRATE-Base (with ReLU activation).
Table 4: Top 1 accuracy of CRATE on ImageNet-1k with different architecture design variants. The soft-thresholding activation  $S_{\lambda \eta}$  is defined in Eq. (117).
Model	MSSA-block	ISTA-block	ImageNet	CIFAR 10*	CIFAR 100*
CRATE-B	default	ReLU activation (default)	70.8%	96.8%	82.7%
CRATE-B	default	soft-thresholding activation	67.6%	96.0%	76.8%
Potential theoretical justification for the performance differential. Previous work [49] studied the phase collapse mechanism for understanding the non-linearities and the convolutional filters used in CNNs such as ResNet [23] on classification tasks. Specifically, they found that replacing the phase collapses with thresholding operators which enforce sparsity largely degrades the classification performance of CNNs. The effect of phase collapse analyzed in [49] is to better separate out the means of different classes within a classification task. This may account in part for the increase in classification accuracy reported in Table 4. On the other hand, we believe that the CRATE architecture will be applicable beyond just classification tasks. In CRATE, the purpose of the training process is to learn the local signal models at each layer (see e.g., Section 2.5). From this perspective, so long as the downstream training task requires semantically meaningful representations of the data distribution, the exact training configuration is of secondary importance. In particular, we may use self-supervised learning methods such as (masked) autoencoding to learn the signal models, whence there may not be any well-defined notion of class mean. In such cases, a priori we may expect both soft thresholding and nonnegative soft thresholding to perform comparably well. We leave the verification of this to future work.
Theoretical justification for non-negative sparse coding. The sparse rate reduction formulation (1) does not include a non-negative constraint, and the token representations have marginal distribution equal to a mixture of zero-mean Gaussians (which are symmetric around 0). Below, we argue that the non-negative sparse rate reduction optimization and the regular sparse rate reduction optimization engender representations which are qualitatively similar in many ways, which confirms our conceptual understanding of how CRATE performs structured representation learning.
First, we formalize the non-negative sparse rate reduction optimization problem. Let  $\chi$  be the characteristic function (with codomain  $\{0,\infty\}$ ) of its input proposition. Then the non-negative analogue to (1) is
$$
\max  _ {f \in \mathcal {F}} [ \Delta R (\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}) - \lambda \| \boldsymbol {Z} \| _ {0} - \chi (\boldsymbol {Z} \geq 0) ] \quad \text {w h e r e} \quad \boldsymbol {Z} = f (\boldsymbol {X}). \tag {118}
$$
Although formal analysis of the optimal points of the sparse rate reduction maximization problem (1) or its nonnegative variant (118) is out of scope of this work, we see that the rate reduction maximization (i.e.,  $\max_{f\in \mathcal{F}}[\Delta R(Z;U_{[K]})]$  has optimal points characterized similarly to [46, Theorem A.6], namely that the representation of each distribution in the mixture is supported on a subspace with nearly isotropic covariance on this subspace, and the supporting subspaces are (nearly) orthogonal. Adding the sparsity term  $\lambda \| Z\| _0$  for some regularizer  $\lambda$  would enforce the axis-alignment of the supporting subspaces; when adding in addition the nonnegativity term  $\chi (Z\geq 0)$ , following through the proof of [46, Theorem A.6] suggests that the argument goes through with suitable modifications (in particular, considering the conclusions for the covariance rather than  $ZZ^{\top}$ ). This sketch suggests that the statistical and geometric properties of the optimal representation remain the same when adding the non-negative constraint to the sparse rate reduction formulation. We leave a detailed proof to future work.
B.6 Pre-training on ImageNet-21K
We investigate a larger pre-training dataset for training CRATE. Specifically, we first pretrain on ImageNet-21K [9], which contains 14 million images, and then fine-tuned on ImageNet-1K. As
shown in Table 5, with the CRATE-Base model (22.80M parameters), we achieve  $80.2\%$  top-1 accuracy; this is comparable to ViT-Base (86M parameters,  $83.9\%$ ) [40] with around  $25\%$  of the parameters. For pre-training on ImageNet-21K, we configure the learning rate to  $1 \times 10^{-4}$ , set the weight decay to 0.05, and use a batch size of 4,096. The total number of epochs is 90, with 10 warmup epochs. For fine-tuning on ImageNet-1K, we use the same set of parameters as described in Appendix B.1.2, with the exception of setting the learning rate to  $5 \times 10^{-5}$  and having a total of 50 epochs.
Table 5: Top 1 accuracy of CRATE-Base and ViT-Base [40] on various datasets when both models are pre-trained on ImageNet-21k. We use the ViT-Base results from [40] as a basis for comparison.
Model	# parameters	ImageNet	CIFAR 10	CIFAR 100
CRATE-Base	22.80M	80.2%	98.3%	88.3%
ViT-Base [40]	86M	83.9%	99.0%	91.7%
B.7 Evaluating  $R^c$  and sparsity for ViT
We conduct experiments to evaluate the  $R^c$  and sparsity of token representations from each layer of a pre-trained ViT-Base (downloaded from https://github.com/huggingface/pytorch-image-models). We summarize the results in Figure 12. We find that without our white-box design, the vanilla ViT does not optimize our proposed sparse rate reduction objective. This contrasts with the results shown in Figures 3 and 4 of the work, wherein we can observe that the compression term  $R^c$  and sparsity value decrease layerwise for CRATE, in accordance with our theory.
Figure 12: Left: The compression term  $R^{c}(Z^{\ell +1 / 2})$  of the multi-head self-attention outputs at different layers. Middle: The sparsity of outputs of the MLP block,  $\| Z^{\ell +1}\| _0 / (d\cdot N)$ , at different layers. Right: To get a more fine-grained understanding of the sparsity of MLP block outputs of ViT, we use three different thresholds  $\tau \in \{1.0,0.5,0.1\}$  and measure  $\sum_{i,j}\mathbf{1}\{|Z_{i,j}^{\ell +1}| < \tau \} /(d\cdot N)$ , where  $Z_{i,j}^{\ell +1}$  represents the  $j$ -th element in the  $i$ -th token representation. (First row model: ViT-Small; second row model: ViT-Base; third row model: our proposed CRATE-Small).