An In-depth Investigation of Sparse Rate Reduction in Transformer-like Models
Authors: Yunzhe Hu
Difan Zou
Dong Xu
Venue: NeurIPS 2024 poster (2024)
Keywords: Sparse Rate Reduction
Transformer
Complexity Measure
Generalization
Primary area: evaluation

Abstract: Deep neural networks have long been criticized for being black-box. To unveil the inner workings of modern neural architectures, a recent work proposed an information-theoretic objective function called Sparse Rate Reduction (SRR) and interpreted its unrolled optimization as a Transformer-like model called Coding Rate Reduction Transformer (CRATE). However, the focus of the study was primarily on the basic implementation, and whether this objective is optimized in practice and its causal relationship to generalization remain elusive. Going beyond this study, we derive different implementations by analyzing layer-wise behaviors of CRATE, both theoretically and empirically. To reveal the predictive power of SRR on generalization, we collect a set of model variants induced by varied implementations and hyperparameters and evaluate SRR as a complexity measure based on its correlation with generalization. Surprisingly, we find out that SRR has a positive correlation coefficient and outperforms other baseline measures, such as path-norm and sharpness-based ones. Furthermore, we show that generalization can be improved using SRR as regularization on benchmark image classification datasets. We hope this paper can shed light on leveraging SRR to design principled models and study their generalization ability.

An In-depth Investigation of Sparse Rate Reduction in Transformer-like Models
Yunzhe Hu
Difan Zou
School of Computing and Data Science
School of Computing and Data Science Institute of Data Science The University of Hong Kong dzouβes hku. hk
The University of Hong Kong yzhu@cs hku. hk
Dong
School of Computing and Data Science The University Hong Kong dongxuβcs hku hk
Abstract
Deep neural networks have long been criticized for being black-box. To unveil the inner workings of modern neural architectures, recent work [45) proposed information-thheortiic objective function called Sparse Rate Reduction (SRR) and interpreted its unrolled optimization Rate Reduction Transformer (CRATE). However, the focus of the study was primarily on the basic implementation, and whether this objective optimized practice and its causal relationship to generalization remain elusive. Going beyond this study, we derive different implementations by analyzing layer-wise behaviors of CRATE, both theoretically and empirically. To reveal the predictive power collect set of model variants induced by varied
Transformer-like model called Coding
SRR on generalization, implementations and hyperparameters and evaluate SRR complexity measure based on correlation with generalization. Surprisingly, we find out that SRR has positive correlation coefficient and outperforms other baseline measures, such as path-norm and sharpness- based ones. Furthermore, we show that generalization can be improved using SRR regularization on benchmark image classification datasets. We hope this paper can shed light on leveraging SRR to design principled models and study their generalization ability.
Introduction
Transformers 39 have become the de facto choice neural architecture nowadays and find great success in applications across language, vision, speech, and other scientific fields. The self-attention module in Transformers utilize global interactions to capture long-range dependency. However, the mechanisms and learning process of self-attention and other components in Transformers remain open problems, calling for more research to interpret and understand their properties.
One approach to interpreting the attention module involves experimental observation of the attention module to gain insights into their behaviors. For instance, DINO (7] provides a means to observe and analyze attention maps w.r.t class tokens in Vision Transformer (ViT), shedding light on their emerging interpretability from self-supervised learning. Another line of work focuses on interpreting or building attention module and even Transformer-like models from mathematical perspective. Works in this vein have attempted to establish connections between Transformers and reverse- eeverse-engineereed energy
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

function 43 17
associative memory such as modern Hopfield network [35 37 26 4 and sparse
distributed memory
[5 or programming languages 41
name a few.
Recently, the study of algorithm unrolling has emerged as a promising technique to bridge the gap between iterative optimization and neural architecture. A work by Yu et al.
45 considers the
objective of representation learning as optimizing the Sparse Rate Reduction (SRR), function that promotes maximum information gain described by the coding rate function 244 46 and induces sparsity. In particular, they show that a Multi-head Subspace Self-Attention (MSSA) operator with skip connection and an Iterative Shrinkage-Thresholddde Thresholding Algorithms (ISTA) operator can be derived under some assumptions by unrolling minimization of the coding rate of representations in incoherent subspaces, i.e., compression and sparse coding, i.e., sparsification, respectively. By stacking these operations into layers, they build Transformer like model CRATE in which every layer should have the completely interpretable compress-then- then-sparssffy behavior. However, although motivated by an information-thhoreee and principled objective, it is still unexplored whether the core component MSSA operator with skip connection indeed implements the idea of compression in practice and how information propagates in the forward pass. On the other hand, SRR the objective of representation learning is still an empirical formulation. Its causal relationship to generalization remains elusive.
In this paper, we conduct an in-depth investigation of this Transformer-like model and take steps to address these limitations. Our contributions are summarized as follows:
In Section 4. we highlight the derivation artifacts through analysis of the key component MSSA operator and explore implementation variants of CRATE by inspecting the layer-wise behaviors. We show that the gradient approximation of the compression term will yield counterproducttve effect, performing decompression of token representations instead.
In Section 5
we uncover the correlation between the learning objective SRR and generalization in unrolled models. By training models with varied hyperparameters, we show that SRR complexity measure has a positive correlation coefficient and outperforms other baselines.
Scction we demonstrate the effectiveness of SRR as a regularization technique for improved performance benchmark datasets. Specifically, we show that the classification accuracy of unrolled models on CIFAR-10/100 can be consistently improved using a simple and efficient implementation of regularization.
2 Related Work
2.1 Interpreting Transformers
has surged recently. Despite its achievements, the mechanisms and learning of attention layers remain enigmatic. One approach to interpreting Transformers is to experimentally observe the inner representations or output of key components like self-attention. This includes analysis by projecting parameters of Transformers to embedding space 10 inspecting the representations with another language model [1], visualizing attention map
Research on interpreting Transformers 13 ||
7 [4]. ctc. Several works opt for "mechanistic interpretability" [12 28 40] aiming to reverse-
engineer the representations learned by Transformers that have "grokked" or mastered complex modular arithmetic task [34] and other synthetic tasks 23 47). Another line of work focuses more on theoretical understanding and building connections to other concepts. These papers utilize tools such Bayesian inference []] convex optimization [36] analyze attention in Transformers. There have also been attempts to interpret a Transformer as an energy function optimizer [3 [7]. connect
attention to memory $55
37
26
4 5
interacting particle systems [3] transform into human-
readable programs 41 222. to name just few. Our work focuses on the empirical investigation of a
Transformer-lis model, CRATE
45
recently introduced from pure mathematical derivation.
2.2 Algorithm Unrolling
Algorithm unrolling [27] has emerged as a promising technique for designing interpretable and efficient deep learning architectures. This approach establishes a direct connection between iterative algorithms and neural architecture, with each iteration of the algorithm corresponding to one layer of the architecture. Previous works have employed this technique to design popular networks in forward-construeee manner. For instance, the seminal work [15] proposed to unroll the Iterative

---

Shrinkage-Thresholddee Thresholding Algorithm for sparse coding into layers of linear operation followed by ReLU non-linearity. Other works have tried find representation objective function to unroll into
convolutional neural network 33 6 graph neural network [42], and Transformers 43 17 We will
follow this iteration-layer correspondence to conduct layer-wise analysis.
3 Revisiting Sparse Rate Reduction
2N] Let 2 = [ E IR dddxN denote N samples, where each column 2; E IRd represents tokens in Transformers. U -U=U....... .Uk......] E RdxKr denote a set of incoherent basis spanning K subspaces, wherein columns of U E RdxP represent basis in i-th low-dimensional subspace (p < d). follow the configuration that d Ky in standard ViT [
Previously, Yu al. 46 propose that the compactness of representations Z E RdxN can be measured
by coding rate function: R Z) log det N z" Z). A more recent study [45 contends that the objective of representation learning is to transform and compress samples from an unknown distribution mixture of low-dimensional Gaussian distributions supported on incoherent bases. This objective boils down to the maximization of Sparse Rate Reduction (SRR):
max Z ZERUXN N
R(2) - R (Z:U) A| lu.
where || ||o means to norm and R = R(U Z) measures the compactness of repre- sentations in the low-dimensional subspaces. One layer of a network, formulated as a mapping Sw (.) parameterized by w. can be interpreted as applying one step of gradient-based methods the objective in (T). In practice, Yu et al. [45] use alternating minimization to break down the optimization into two steps: compression, i.c. minz R (Z; U) and Sparsification, i.c. minz ^ll lo R(Z). Specifically, given representation z--1 at (E 1)-th layer, Z' can be obtained by two-step optimization:
Y 2 -1 a. R Z
U ~ z f-1 aa MSSA(Z" 8- U²)
(2)
z ReLU - -x11)
where a,33 are step sizes, D E Raxd is assumed as a complete dictionary, scalar y = Nea and
K
MSSA(Z:U) Σ UR (00)
U Z softmax (UT
(4)
U Ukl
U Z (Uft
The operator MSSA(:U) in called the Multi-head Subspace Self-Attention (MSSA) operator, takes the form of self-attention in standard Transformers (39 [1] with tied query, key and value matrix, i.e., U while the output matrix being its transpose, i.e. UR. Instead of strictly following this formulation, they further replace [U....., Us] E kdd in the MSSA operator with an additional learnable parameter W E R4x Kp. To distinguish them, we name the model with implementation 44) CRATE-Clonccptuall minimization, a Transformer-like model with layered structures can be naturally constructed. With input Z", c.g. tokenized images in ViT, an L-layer model iteratively optimizes the input and yields the final representations zh_ Parameters ( and (D*) 1/ can be learned through end-to-end training [15].
By incrementally optimizing (I) with alternating
4 Is Sparse Rate Reduction Optimized in Transformer-like Models?
While the white-box Transformer-like model proposed 455 is derived by unrolling optimization upon pre-defined objective function, whether the optimization is implemented by the model in the forward pass still unclear. In this section, we first review the main derivations at the core of building CRATE, i.e. unrolling optimization minz R (Z: U)) into MSSA operator with skip connection as in ( and identify the pitfalls in implementing the minimization. We then provide variant models based on different implementations and empirically show their wise behaviors.

---

2.0
R9(Z-))U0
log
1800
R922409
1.5
15
1600
1.0
112
R";;
Output
1400
1200
0.5
-1.0
1
2
5
6
9
10
11
12
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4,0
Layer
(a) Changes of coding projected the same set basis at each layer.
(b) Different approximation of log(-).
Figure 1: In simplified attention-only experiment, MSSA operator with skip connection actually implements an ascent method on (2; U). opposed to its design purpose (left). This due to an artifact approximation with its second-order term. (right)
4.1 Pitfalls in Deriving CRATE-C
first show that second-order Taylor expansion of the coding rate representations projected onto subspaces can be expressed as:
X - (A - 1)²
RC RZ(ZU))
log
k=1 im
(5)
First-order term
Sccond-order term
where X 1,i [N] are the eigenvalues of + +U z U Z. Following the derivation and implementation from Appendix A.2 in [45. the MSSA operator with skip connection is constructed by performing an approximation of gradient descent on R (Z:U):
K
- (= oz Σ U ( (1+)0?Z)
(6)
K JUUTZ ~ Euuu
2 Z a
(7)
Vof first-order
Va second- order
Z + 0r Σ 2 2U
(8)
can be seen that this update step takes the gradient of a lower bound of R [Z; U) and discards the first-order term. With a proper step size, the coding rate the same subspaces is expected decrease after one iteration. However, we will show that this is not the actual case via a toy experiment.
We consider a simplified setting where layers update 8 are conducted with parameters {U') initialized as orthonormal matrices. We initialize a random variable Z from a Gaussian distribution and measure the coding rate before and after each layer. We set N 196. L 12, d = 384, K = 6. 1, and proper e? such that y 1. As shown in Figure la R (Z'; U³) is always greater than R ( U³) and R pr"": U²) is increasing in general the layer goes deeper. This means the update implement descent method R. The crux lies in the approximation of Re's gradient.
( that resembles the standard self-attention with skip connection does not essentially
When taking the gradient of Re to construct the MSSA operator, omitting its first-order term will produce counterproductive effect. As shown on the left-hand side of the inequality in (5). R can be

---

expressed as the sum of logarithms of eigenvalues. We expect the eigenvalues to decrease to minimize the value of R Figure Tb iillustaates different approximations of the logarithm function. If we omit the first-order term of its Taylor expansion and only perform descent methods on its second-order
term (corresponding to (A' - the eigenvalues will go up leading to an increase in the value of R.. Therefore, one step of update (8) secretly maximizes R², contrary to the purpose of its design.
More figures detailing this issue are in Appendix A
4.2 Producing CRATE Variants
In the previous subsection, we show the problems arising from gradient approximation when unrolling minz (2; U) into MSSA operator with shorteut. We will, in the subsection, introduce two variants of CRATE induced by the conceptual and implementation gaps. These variants can be considered as the alternative instantiations of the optimization geimizzaionninducedd architectures but in a more self-contained way. They also serve representative samples for our subsequent investigations of SRR.
One variant of CRATE, motivated by the theoretical gap between CRATE-C and the SRR principle, could naturally emerge when the sign before the MSSA operator in (8) is changed. Similar to previous analysis via eigenvalues, this update of representations in fact implements one step of ascent methods on the second-order term of Rt therefore minimizing the eigenvalues and consequently R This variant is designed to counter the pitfalls in CRATE-C, enabling a more faithful reduction in R2 and thereby enhancing alignment with the SRR principle. We term the Transformer-like model with this implementation CRATE-N(egative::
Z - ar? Σ
The other variant we would like to introduce is motivated by the misalignment between CRATE and CRATE-C. Although replacing the output matrix [U....., Ux] with learnable parameters W in CRATE empirically boosts performance, it also contaminates the framework and sacrifices the mathematical interpretability. Does this modification really matter? Can we preserve model performance while maintaining framework integrity? turns out that simple transpose operation of the output matrix could greatly close the empirical gap to CRATE, without more parameters. Other manipulations and discussions can be found in Appendix B We refer to the model with this simple manipulation CRATE-T(ranspuu
Uf Z softmax (10
+ an UU,,,, UK)²
(10)
Uk Z (aniiimax(UT,Z)
4.3 Behaviors of Sparse Rate Reduction
The Transformer-like model CRATE is built by sequentially stacking the layer that comprises two modules 2 and (3) (or different implementations). Although each module is designed to implement one-step optimization of different objectives, it is unclear whether the architecture design achieves the optimization as a whole. On the other hand, there also a need to determine whether the model parameters learned through end-to-end training actually lead to improved optimization.
To investigate how sparse rate reduction evolves in the forward pass and during training, we train CRATE and its variants on CIFAR- 10/100 datasets and evaluate the sparse rate reduction measure R" at different layers and epochs on the training set. À is chosen as 0.1 and detailed experiment settings can be found in Section 6. Figure 2 and Figure 3 show the behaviors of sparse rate reduction of CRATE along with its variants CRATE-C, CRATE-N, and CRATE-T under Tiny configurations in 45). When the models are randomly initialized, the sparse rate reduction measure almost monotonically decreases in the first 9 layers and then rises in the subsequent layers. This partly confirms the layer-wise optimization of the objective SRR and its alignment with forward architecture design, although in Section we demonstrate that Rº U) will monotonically go up the absence of operation B3. We conjecture that the ReL.U non-linearity may also play an important role in optimizing the compression term R [Z; U) in the forward pass.
A| + R R ;
5

---

75
3750
3500
3400
3250
3200
3000
3000
2750
2250
2400
2000
10
11
10
12
Layer
Layer
(a) Sparse rate reduction measure of CRATE-C 8
(b) Sparse rate reduction measure of CRATE-N N
Eaach
Spach 123
36.00
Epocch 73
2pech
Epoch 1poch
baach
3600
Esoch spech 20
Spach 200
2
3400
3200
3000
3000
2600
2600
2400
2400
Layar
Layar
(c) Sparse rate reduction measure of CRATE-T (T0
(d) Sparse rate reduction measure of CRATE
Figure 2:
Sparse rate reduction measure Z (Z: R(Z) of CRATE and its variants
evaluated at different layers and epochs on CIFAR-10.
Another surprising finding is that as the learning process proceeds, the sparse rate reduction measure each layer will increase monotonically across all models, with exception in the last few layers of CRATE-C
These phenomena give us implications for understanding Transformer-like like models: the representations of initialized models converge fast in the first few layers and hover around the local minimum of the objective landscape; however, the useful information in representations may be discarded due to over-compression and the learning of parameters gradually increases sparse rate reduction measure to counteract this effect for improved task-specific representations.
To summarize, our finding is that sparse rate reduction measure is incrementally optimized realistic setting at initialization. This aligns well with its design purpose from a macro perspective. With varied implementations, the result still holds even when the compression-inppins MSSA diverges from its goal from micro perspective. We postulate that ReLU non-linearity in (3) could also promote compression and leave their interaction for future work.
operator
Whether Sparse Rate Reduction Benefits Generalization?
So far, we have partially confirmed the validity of different implementations of Transformer-lh models by inspecting the layer-wise optimization of SRR. But whether this objective important or principled for these architectures to generalize is still an unaddressed problem. In this section, we want to explore the predictive power of SRR and its causal relationship the generalization of CRATE.
5.1 Sparse Rate Reduction Complexity Measure
An important tool to study the generalization of deep networks is complexity measure. A complexity measure that can properly reflect the generalization needs to have the following property: lower complexity should indicate smaller generalization gap. Complexity measures can be either theoreti-
cally motivated, such as PAC-Bayes [25] 29], VC-dimension (8), norm-based bounds [3 3 30 or
empirically motivated, such as sharpness 20 and path-norm
31. We choose to adapt SRR into

---

Eanth 25
153
Epoch
211.. (2 R121
3600
4000
3500
3200
3000
2500
2600
2400
2000
10
10
12
Layer
Layer
(a) Sparse rate reduction measure of CRATE-C C 8
(b) Sparse rate reduction measure of CRATE-N N
4000
Sp
Spuch 250
2
4000
Spech spech 20
Epacc 73
105
2pech
seach
20
325
Spech 200
3750
Spach 200
3500
X200
J000
3000
2800
2750
2500
Layar
Layar
(c) Sparse rate reduction measure of CRATE-T (10
(d) Sparse rate reduction measure of CRATE
Figure 3: Sparse rate reduction measure A| Z (Z: R(Z) of CRATE and its variants evaluated at different layers and epochs on CIFAR-100.
complexity measure that belongs to the latter category:
MSRR (SRR(W:; Z) 1 Σ t=1
HSRR (w:: zz Σ ( R(())
(11)
where z denotes the output at layer é and w contains the parameters including U and Dº.
5.2 Correlation with Generalization
An effective measure of complexity should bound the generalization gap, defined as the difference between validation loss and training loss when the latter reaches threshold, i.e., L Leraiin with high probability, However, for those measures that do not provably bound this gap, as is the case with SRR measure (II), we need evaluate its correlation with the generalization gap understand its causal relationship to generalization.
Collecting Trained Models To evaluate the complexity measure and generalization across models, consider changing the hyperparameters and collect set of models trained to meet specific stopping criterion. Here, we also consider the model type containing different variants of CRATE hyperparameter investigate its influence generalization. Formally, let O denote type of hyperparameter with |O,] different choices, and define (01, o...... 8.) E θ x ... θ, as an instantiation from types of hyperparameters. By varying choices across hyperparameter space, we can produce [01] x ....X x (e,.] models. In our experiment, we consider n = 5 hyperparameters, including batch size, initial learning rate, width, dropout, and model type. Each contains 2 choices except that the model type contains 4 implementations we discussed before. We successfully train a total of 64 models CIFAR-10 dataset, when cross-entropy loss reaches 0.01 following the stopping criterion in [8] Experimental details and choices of hyperparameters can be found in Appendix [C
common method for measuring correlation by utilizing Kendall's rank correlation coefficient [19 18] which ranges from -1 to 1. Generally, the closer the coefficient is to one, the stronger the causal relationship and the greater the predictive power a measure can offer for generalization. Zero value usually means independent relationships. For given complexity
Evaluation Criterion
7

---

Table I: Correlation of complexity measures with generalization gap (width d 384).
Complexity measures
Batch size
Learning rate
Dropout
Model type
Overall
y
f2-norm
0.200
-0.333
-0.333
-0.429
-0.363
-0.224
norm-iiii
0.200
-0.200
-0.333
-0.286
-0.290
-0.158
# params
0.000
0.000
0.000
-0.572
-0.351
-0.143
1/margin
-0.067
0.467
0.467
0.238
0.415
0.276
sum-of-spec
0,200
-0.333
-0.467
-0.381
-0.290
-0.245
prod-of-spce
0,200
-0.333
-0.467
-0.476
-0.338
-0.269
sun-of-spec/marrin
0.333
-0.333
-0.467
-0.048
-0.230
-0.129
wrod-of-spec/margee
0.333
-0,333
-0.467
-0.143
-0.260
-0.152
fro/spec
-0,200
0,333
0.467
-0.776
0.019
0.031
spec-init-main
0.333
-0.333
-0.467
-0.190
-0.273
-0.164
spec-orig-mate
0.200
-0.333
-0.467
-0.095
-0.252
-0.174
sum-of-fro
0.200
-0.333
-0.333
-0.381
-0.325
-0.212
prod-of-fro
0,200
-0.333
-0.333
-0.429
-0.372
-0.224
sum-of-Fi cum-of-fro/marginn froommrgiin
0.333
-0,200
-0.467
-0.048
-0.217
-0.095
0.333
-0,200
-0.477
-0.143
-0.247
-0.119
fro-distance
0,200
-0,200
-0.333
-0.286
-0.290
-0.155
spec-distance
0,200
-0.200
-0.333
-0.286
-0.290
-0.155
param-norm
0,200
-0.333
-0.333
-0.429
-0.363
-0.224
path-norm
0.333
-0.600
-0.467
-0.286
-0.191
-0.255
pac-bayes-init
0.200
0.200
-0.600
0.238
0.015
-0.009
bac-bayes-orig
-0.200
0.333
0.467
0.381
0.333
0.245
1/ pac-bayes-Slaee flatness
-0.267
0,333
0,333
0.455
0.333
0.213
SRR
-0.067
0.467
0.333
0.714
0.445
0.362
measure, we can construct set of samples containing the measure (0) and generalization gap (0) evaluated different combinations of hyperparameters and calculated Kendall's coefficient on this set:
Uoce x x.. (H(())
(12)
(T) [T(TT
Σ
Σ
sigu (Mz #2 sign (01 92)
(13)
(p1-9)ET (Pr.)))77MP..0
Experimental Results. In our experiment, we find that the correlation of various measures with the generation can be reflected with more prominence under a selected width. Accordingly, present the results terms of Kendall's coefficient r in Table and scatter plot of SRR measure Figure when the width d is chosen 384. The results when granulated coefficient V is also reported (see [18] for detailed definition).
768 is deferred Appendix D The
We confirm the findings from prior works that some norm-based measures, such as sum/prod of spectral/Frobenius norm of parameters neg- atively correlate with generalization, even on Transformer-like models. An interesting find- ing is that path-norm also negatively correlates with generalization, which partly contradicts the previous conclusion. This implies that regular-
Correlation of SRR Measure (T = 0.445)
3.5
CRATE-C
CRATE-N
3.0
Generalization
CRATE-T
CRATE
2.5
ization on path-norm, e.g. Path-SGD (3I), may not be applicable for improved generalization on Transformer-like models. Among the mea-
2.0
1.5
sures we investigated, the inverse of margin and sharpness- based PAC-Bayes flatness show posi- tive and strong correlations. This result justifies common belief that larger margin or flat- ter loss landscape leads to better generalization across the investigated Transformer-like models. Compared baselines, the SRR measure in
0.230 0.240 0.245 0.250 Value Complexity Measure
Figure 4: A scatter plot illustrating the value of SRR measure and generalization gap across CRATE and variants with network width d = 384.

---

achieves the highest overall coefficient and, particularly in the model type axis, outperforms the rest. This motivates the use of SRR regularization the loss function to improve generalization.
Sparse Rate Reduction as Regularization
Since SRR measure enjoys a strong correlation to the generalization of Transformer-lin models, we would like investigate its potential the direct regularization the standard training loss. In particular, we add the SRR measure in IT by regularization factor 77 to the cross-entropy loss:
min (w) + A
1 L
Σ
(w"; Z StopGraa)
(14)
where X > O is the regularization coefficient and StopGrad - Sm (StopGrad The operator StopGrad here, implemented as "Tensor. detach()" in PyTorch, prevents gradient propagation from the output Z' to the previous layers. This allows parameters we at each layer to be updated without interfering with each other, giving more precise optimization of SRR in separate layers.
6.1 Experiment Settings
Model Configurations We follow the configuration of CRATE-Tiny in [45] in this experiment. Specifically, we set the depth L = 12, width d = 384, number of subspaces K = 6, step size = and scaling factor 7 = We also include LayerNorm before each operation in (3 for better trainability and learnable positional encoding. A trainable [CLS] token is prepended to the representations for computing cross-entropy loss and classification.
Datasets and Optimization We use CIFAR-IO and CIFAR-100 datasets for training and evaluation. In practice, we adopt Adam (21] optimizer and initialize learning rate as 10-4 with cosine decay. All models trained for 200 epochs with batch size 128. Note that we only the basic data augmentations: random resize and cropping, horizontal flipping, and RandAugment [9] (with and magnitude = 14). We do not
the number transformations state-of-the-ar performance but to demonstrate the effectiveness of SRR as regularization. We tune the factor via grid search over f0.0001, 0.001, 0.01, 0.1, 1} and find that 0.001 works best. All experiments are conducted on NVIDIA GeForce RTX 3090.
other techniques for
Table 2: Top-1 accuracy for CRATE and its variants trained with or without SRR regularization on CIFAR-10/100 from scratch (width = 384).
CIFAR 10
CIFAR-100
Models
cross-entropy
+ SRR regularization (L=12)
cross-entropy
SRR regularization (L=12)
CRATE-C
76.87
77.61
43.40
44.53
CRATE-N
81.52
81.91
$5.11
55.62
CRATE-T
85,49
85.52
60.59
60.69
CRATE 86.67
86.79
62.40
62.52
6.2 Efficient Implementation
Regularizing the training loss with sparse rate reduction measure [T needs to compute R(Z) and R (Z:U) for every layer. However, this is highly inefficient as it involves high- siggh-imensional matrix multiplication, and it lacks flexibility in controlling parameters. To alleviate this issue, we implement efficient regularization as per layer regularization or random layer regularization: select pre-defined layer or a random layer with uniform probability during training. In practice, we find that former works better. Table provides the results of CRATE and its variants trained from scratch on CIFAR-10/100. SRR regularization is sufficient to improve the performance by simply leveraging the last layer. We also provide a comparison of efficient implementations in Appendix

---

7 Conclusion
To further research in interpreting neural architecture, we provide an in-depth investigation of a recent mathematically driven Transformer-Hn like model, CRATE. Although designed with a principled objective, we identify an artifact in its forward construction and show that the simplest implementation can have the opposite effect in realizing its designed goal. We then provide implementation variants and investigate their layer-wise behaviors in optimizing SRR. An interesting finding is that alternative models exhibit similar behaviors, validating the use of SRR in designing Transformer-lic models. Furthermore, we demonstrate its positive correlation to generalization and effectiveness over baselines. Driven by this connection, we show simple way use SRR as regularization improve performance on CIFAR-101 10/100 datasets. Future direction may include applying layer-wise training and connecting SRR the Forward- Forward algorithm [66] or exploring the impact of depth in the unrolled models.
Limitations
This study has several limitations. Firstly, the conclusion that the SRR measure can be strong indicator of generalization is limited to the CRATE family. Generalizing this conclusion standard Transformers would be non-trivial, as the SRR measure is not properly defined when the query- key-value matrices have independent learnable parameters instead of shared ones. Secondly, the performance of a more faithful implementation (CRATE-N) falls behind the one with simple manipulation (CRATE-T). This calls for rigorous inspection of each component's 's functionality in the framework. Lastly, while confirm the positive correlation to generalization, our analysis is limited in scale. Consequently, drawing definitive conclusions regarding whether SRR can principle or necessitates further engineering to push the model's limit is challenging. A Better and more systematic way needed to determine whether SRR is principled for designing the Transformer- like models and quantify this relationship in appropriate task, perhaps beyond classification.
Acknowledgments
This work was supported in part by Natural Science Fund China (62306252), in part by the Hong Kong Research Grants Council General Research Fund (17203023), in part by the Hong Kong Research Grants Council Collaborative Research Fund (C5052-23G), Jockey Club Charities Trust under Grant 2022-0174, in part by the Startup Fund and the Seed Fund for Basic Research for New Staff from The University of Hong Kong, and in part by the funding from UBTECH Robotics.
part by The Hong Kong
References
[1] Bang An, Jie Lyu, Zhenyi Wang, Chunyuan Li, Changwei Hu, Fei Tan, Ruiyi Zhang, Yifan Hu, and Changyou Chen. Repulsive attention: Rethinking multi-head attention bayesian inference. arXiv preprint mmXiv:2009.09366 2020.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1rXiv:1607.06450 2016.
[3] Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Spectrally-normalii margin bounds for neural networks. Advances in neural information processing systems, 30, 2017.
[4] Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: memory viewpoint. Advances in Neural Information Processing Systems, 36, 2023.
[5] Trenton Bricken and Cengiz Pehlevan. Attention approximates sparse distributed memory. Advances in Neurul Information Processing Systems, 14:15301-5 2021.
[6] Joan Bruna and Stéphane Mallat. Invariant scattering convolution networks. IEEE transactions on pattern analysis and machine intelligence, 187211887) 2013,
10

---

[7] Mathilde Caron, Hugo Touvron. Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self caf-ssueerviee vision transformers. In Proceedings of the IEEF/CVF international conference on computer vision, pages 9650-9660, 2021.
[8] Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 782-791, 2021.
[9] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition workshops, pages 702-703, 2020.
[10] Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant. Analyzing transformers in embedding space. arXiv preprint arXiv:2209.02533, 2022.
[11] 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. In International Conference on Learning Representations, 2020.
[12] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https//ttansforme-ccrciiii putpeerrrrriiuiiusuuu2222222222 /frameworrkindex..nne htmll.
[13] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2023.
[14] Asma Ghandeharioun, Avi Caciularu, Adam Pearce, Lucas Dixon, and Mor Geva. Patchscope: unifying framework for inspecting hidden representations of language models. arXiv preprint arXiv:2401.061000 2024.
[15] 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.
[16] Geoffrey Hinton. The forward-forward algorithm: Some preliminary investigations. arXiv preprint arXiv:2212.133455 2022.
[17] Benjamin Hoover, Yuchen Liang. Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Horng Chau, Mohammed Zaki, and Dmitry Krotov. Energy transformer. Advances in Neural Informa- tion Processing Systems, 36, 2023.
[18] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv:1912.02117. 2019.
[19] Maurice G Kendall. A new measure of rank correlation. Biometrika, 30(1/2):81 1938.
[20] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprini arXiv: (6000048366 2016.
[21] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
[22] David Lindner, János Kramár, Sehastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Mikulik. Tracr: Compiled transformers as a laboratory for interpretability. Advances in Neural Information Processing Systems, 36, 2023.
[23] Ziming Liu, Eric Michaud, and Max Tegmark, Omnigrok: Grokking beyond algorithmic data.
In The Eleventh International Conference Learning Representations, 2022.
11

---

[24] Yi Ma, Harm Derksen, Wei Hong, and John Wright. Segmentation of multivariate mixed data via lossy data coding and compression. IEEE transactions on pattern analysis and machine intelligence, 29(9):1546-11 1546-1562, 2007.
[25] David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learning theory, pages 164-170, 1999.
[26] Beren Millidge, Tommaso Salvatori, Yuhang Song, Thomas Lukasiewicz, and Rafal Bogacz. Universal hopfield networks: general framework for single-shot associative memory models. In International Conference on Machine Learning, pages 15561-15583. PMLR, 2022.
[27] Vishal Monga, Yuelong Li, and Yonina C Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38(2):18 44, 2021.
[28] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretabilityy. arXiv preprint arXiv:2301.052117 2023.
[29] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems, 30, 2017.
[30] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrally-normaizzed margin bounds for neural networks. arXiv preprint arXiv:1707.0955 2017.
[31] Behnam Neyshabur, Russ Salakhutdinov, and Nati Srebro. Path-sgd: Path Rath-normalized optimization in deep neural networks. Advances in neural information processing systems, 28, 2015.
[32] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Norm-based capacity control in neural networks. Conference on learning theory, pages 1376-1401. PMLR, 2015.
[33] Vardan Papyan, Yaniv Romano, Jeremias Sulam, and Michael Elad. Theoretical foundations of deep learning via sparse representations: A multilayer sparse model and its connection to convolutional neural networks. IEEE Signal Processing Magazine, 35(4):72-89, 2018.
[34] Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- cralization beyond overfitting on small algorithmic datasets. arXiv preprint arXiv:2201.02117, 2022.
[35] Hubert Ramsauer, Bernhard Schäfl. Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need, arXiv preprint arXiv:2008.022: 2020.
[36] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In International Conference on Machine Learning, pages 19050-19088. PMLR, 2022.
[37] Danil Tyulmankov, Ching Fang, Annapurna Vadaparty, and Guangyu Robert Yang. Biological learning in key-value memory networks. Advances in Neural Information Processing Systems, 34:22247-22258, 2021.
[38] Vladimir N Vapnik and A Ya Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. In Measures of complexity: festschrift for alexey chervonenkis, pages 11-30. Springer, 2015.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[40] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability the wild: eircuit for indirect object identification in gpt-2 small, arXiv preprint arXiv:2211.00593 2022.
12

---

[41] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pages 11000-10900. PMLR, 2021.
[42] Yongyi Yang, Tang Liu, Yangkun Wang. Jinjing Zhou, Quan Gan, Zhewei Wei, Zheng Zhang, Zengfeng Huang, and David Wipf. Graph neural networks inspired by classical iterative algorithms. In International Conference on Machine Learning, pages 11773-11783, PMLR, 2021.
[43] Yongyi Yang, David P Wipf, al. Transformers from optimization perspective. Advances in Neural Information Processing Systems, 35:36958-36971, 2022.
[44] Catherine Ych, Yida Chen, Aoyu Wu, Cynthia Chen, Fernanda Viégas, and Martin Wattenberg. Attentionviz: A global view of transformer attention. IEEE Transactions Visualization and Computer Graphics, 2023.
[45] 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.
[46] Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. Learning diverse and discriminative representations via the principle of maximal coding rate reduction. Advances in Neural Information Processing Systems, 33:9422-999 0039422-2344. 2020.
[47] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task, arXiv preprint arXiv:2206.0430. 2022.
13

---

A Complete Demonstrations of the Pitfalls
To give clearer picture of how approximations affect the optimization of R we provide the complete results with different update rules under the same simplified settings in the main text:
(a) Gradient descent on R
(b) Gradient descent the second-order Taylor expansion of R
(c) Gradient descent on the first-order term of the Taylor expansion of R (w/o second term).
(d) Gradient descent the second-order term of the Taylor expansion of R (wlo first-order term).
(e) Further adding softmax function upon (d).
The results in Figure [ correspond c to the above experiments. Gradient descent on Rt did make it decrease across layers. Conversely, applying gradient descent on its second-order Taylor expansion resulted in an increase, indicating potentially flawed approximation. Isolating gradient descent to the second-order term led to rise in R as opposed to the design purpose. Furthermore, incorporating the softmax function, real-world operation examined in the main text, did not alter this conclusion.
A921-100
A-22--1,10
592
mizi.09
642
1332
523
1230
750
700
1030
602
(a)
(b)
(c)
1500
1400
1605
1100
1000
31
Layer
(d)
(e)
Figure 5: (a) Original gradient update, i.e (6 (b) update from second-order Taylor expansion, i.c., (c) update from removing the second-order term from (7) (d) update from removing the first-order term from (e) update from further adding softmax, i.e.,
B Different Manipulations to the Output Matrix
the MSSA As mentioned in Section CRATE replaces the output matrix (U...., operator with learnable W (which is different from U). We then raise the following question on the manipulation of the output matrix: if we are free to adjust the output matrix while sacrificing interpretability, can find more alternatives that can outperform CRATE-C or even CRATE? In practice, we have experimented with setting this matrix to an identity or fixed randomly initialized matrix, but only to discover that transpose performs best (Table 3). Therefore, CRATE-T is a feasible choice without introducing new parameters, which can be utilized to better understand the SRR principle and its connection to the performance.
want to clarify that the analysis here intends to compare the variants with CRATE-C. not CRATE, because CRATE introduces learnable parameters that are less interpretable. We believe there are at least some interesting conclusions from the comparison: 1) CRATE-N achieves better performance by following the SRR principle more faithfully, shedding light on the connection of
14

---

Table 3: Top-1 accuracy for CRATE and its variants trained on CIFAR-10 from scratch (width
d = 384).
Models
CRATE-C
CRATE-N
CRATE-T
CRATE
CRATE-Fix
CRATE-Identity
# Params 3.94M 76.87
3,94M 81.52
3,94M 85.49
5.7IM 86.67
3.94M 80.73
3.94M
Accuracy
83.18
SRR to generalization; 2) We need to explore more design choices (e.g.. CRATE-T, which may deviate from directly optimizing the SRR but still exhibit a similar architecture) to gain a complete understanding of the SRR principle for model performance (this motivates our Section
C Experimental Details of Collecting Trained Models
Our experimental details to generate a family of trained models largely follow the previous work [18] Models with heavy data augmentations tend to generalize better than those without them. is therefore crucial to isolate the influence of data augmentations from the change of other perppaarametrs.. We choose to remove data augmentations during training to ensure that most models can be trained meet the stopping criterion. We include Layer Normalization [2] before each operator during training, but also remove it when evaluating the complexity measures.
In this experiment, we vary across width, dropout probability, and model type. We present the choices of these hyperparameters in Table Adam [2] is used as the default optimizer. Model depth is kept as L 12 and number of subspaces K = 6. Dropout is applied after adding positional encoding, softmax function, and output projection MSSA operator.
of hyperparameters, i.e., batch size, initial learning rate,
Table 4: Choices of hyperparameters.
Hyperparameters
Choices
batch size initial learning rate width
[64, 128]
{2 10 1 10-4] (384, 768)
dropout model type
(0.0, 0.1]
(CRATE-C, CRATE-N, CRATE)
D Correlation of Complexity Measures when width d 768
Table 5 and Figure 6 give results on correlation to generalization when width d = 768. We see that SRR slightly better than other baseline measures in terms overall r. In the axes of dropout and model type, however, it underperforms PAC- PAC-Bayes flatness measure. This implies that width could have a considerable influence on studying SRR as a complexity measure. We leave it for future work.
E
Comparisons of Efficient Implementations
Table compares different efficient implementations of SRR regularization. We find that randomly choosing layers to regularize generally worsens the performance. While regularizing shallower layers may bring more performance gain, leveraging the last layer already suffices outperform the cross-entropy baseline. Specifying which layer to regularize could be expensive, especially when the model size grows. We opt for the last layer, which should be reasonable if depth scales. Our results indicate that this intuitive choice can already give consistent performance gains in different settings.
15

---

Table 5: Correlation of complexity measures with generalization gap (width d = 768).
Batch size
Learning rate
Dropout
Model type
Overall 7
V/
f2-norm
0.000
-0.375
-0.625
-0.250
-0.310
-0.313
fn-normiinit
0.000
-0.375
-0.625
-0.208
-0.274
-0.302
# params
0.000
0.000
0.000
-0.295
-0.188
-0.074
1/margin
-0.125
0.375
0.625
-0.208
0.173
0.167
sum-of-spec
0,000
-0.375
-0.625
-0.755
-0.310
-0.344
prod-of-spec
0.000
-0375
-0.625
-0.417
-0.339
-0.354
sum-of- sum-of-speccmarrin
0.000
-0375
-0.625
00.58
-0.119
-0.655
prod-of-spec/margin
0.000
00375
-0.255
00.17
-03277
-0.354
fro/spec
0.000
0.375
0.500
0.083
0.242
0.239
spec-init-man inne-iiit-mmnn
0.000
00375
0.625
00.177
-0.311
-0.544
spec-orig-
0.000
-03755
-0.255
00.417
-0.331
-0.354
sum-of-fro
0.000
00.755
0.625
00.333
-03066
-03333
prod-of-fro
0.000
-03775
0.625
0.250
-0.278
-0.313
sum-of ofo-oofffff//rinnnnnn
-0.125
00375
0,500
-01677
-02866
-0,922
mrod-of-fro/marre
-0.125
0.375
0.500
00.155
-02388
-0.281
fro-distance
0,000
-0.755
0.625
00208
-0274
-0.302
spec-distance
0.000
-0.375
0.625
00.417
+0322
-0.354
param-norm
0.000
-0.375
-0.255
-0.500
-0.310
-0.316
path-norm
-0.250
-0.625
0.125
-05500
-0.415
-0.313
pac-bayes-in -iit
0.000
-0.375
-0.625
0,250
-0.214
-0.188
pac-bayes-orig
0.000
0.375
0.625
0.167
0.315
0.292
1/ pac-bayes-flatness
0.000
0.375
0.688
0.573
0.337
0.409
SRR
0.125
0.500
0,250
0.375
0.407
0.313
Correlation of SRR Measure (T = 0.407)
CRATE-C
2.6
CRATE-N
Generalization Gap
2.4
CRATE-T
CRATE
2.2
2.0
1.8
1.6
1.4
0.29
0.30
0.31
0.32
0.33
0.34
0.35
Value of Complexity Measure
Figure 6: A scatter plot illustrating the value of SRR measure and generalization gap across CRATE and variants with network width d 768.
16

---

Table Top-1 accuracy for CRATE and its variants trained with efficient implementations of SRR regularization on CIFAR- 10 from scratch (width = 384).
CIFAR-10
Training methods
CRATE-C
CRATE-N
CRATE-T
CRATE
cross-entropy
76.87
81.52
85.49
86.67
Layer reg
77.75
82.41
85.84
87.03
Layer reg
77.95
81.57
85.46
87.03
Layer reg
77.48
80.83
85.22
87.02
Layer reg
77.04
81.29
85.12
86.64
Layer 10 reg
77.44
81.19
85.68
86.67
Layer 12 reg
77.61
81.91
85.52
86.79
Random layer reg
75.19
79.66
84.27
85.36
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We list our contributions in the introduction derived from the main sections.
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 the paper and important assumptions and limitations. A No or
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 can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We provide an additional section discussing the limitations of our work in the manuscript.
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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
17

---

The authors should discuss the computational efficiency of 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, 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 a complete (and correct) proof?
Answer: [NA]
Justification: Our paper is mostly about empirical evaluation and therefore does not provide novel theories.
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 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.
Thcorems 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 the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We faithfully describe our experimental details in the main text.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer 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 make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
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
18

---

(a) If the contribution is primarily a new algorithm, the paper should make it 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 access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to deseribe the particular way they provide for reproducibility. In the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer:
[No]
Justification:
Our paper builds heavily previous works which are publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
mublic/guidess http//gie//ooeeduiiiiieeeseeee
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for 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
/ //nips. www.c/pwwwp..pcdccoocooooo http://gdesssseeeee http//oeeeuaiiisioooooooooooee
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. 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).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: 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?
Answer: [Yes]
Justification:
All the detailed are stated in the main text.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
19

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Due to lack of time, do not include error bars in experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but 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 (c.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 execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: use Nvidia 3090 GPU for all the experiments. Guidelines:
The answer NA means that the paper does include experiments.
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 of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:p///...aapes..cccpppaataats
Answer: [Yes]
Justification:
We strictly follow the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS 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).
20

---

10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: We believe this work makes preliminary efforts towards understanding potentially principled way of designing Transformer-ike models. But there are still gaps and limitations on the implementations of sparse rate reduction as a guide to model design. The conclusion of our paper is still not definitive enough to make strong societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
This is no such risk.
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 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 we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators 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?
Answer: [Yes]
21

---

Justification:
We have properly cited the most relevant papers which we build our code on.
Guidelines:
The answer NA means that the paper does not use 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, if 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 (e.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 use in the package should be provided. For popular datasets, paperssiippssiiii thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 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: There such assets introduced in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of 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 as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, 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
Question: 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?
22

---

Answer:
[NA]
Justification:
The paper does not involve crowdsourcing nor 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 the institution conducting the review.
23


Reviews:
[{"review_id": "xMR0xJn3Lw", "reviewer": "Reviewer_RBPW", "summary": "This paper conducts an in-depth of study of CRATE, a previously proposed Transformer-like architecture to make deep learning more white-box. CRATE was motivated by sparse rate reduction (SRR), and it is a multi-layer architecture designed to optimize the SRR objective iteratively layer by layer. The authors first study whether the CRATE architecture can really optimize the SRR. They empirically show (in Figure 1) that the SRR objective actually increases with layer through a toy experiment, and the reason is that CRATE discards the first-order term in the Taylor expansion of the SRR objective. The authors then propose two variants, CRATE-N and CRATE-T, and study their behavior with experiments. Finally, the authors propose to use SRR as a complexity measure for predicting generalization. They experiment with a bunch of CRATE models and find that SRR correlates the best with the generalization performance. Motivated by this, the authors study whether SRR is a good regularizer.", "questions": "1. I cannot see why CRATE-N makes sense. If it is gradient ascent, how does it minimize $R^c$?\n2. In line 170 the authors claimed that SRR \"rises slightly in the subsequent layers\". However, looking at Figures 2 and 3, my feeling is that it rises significantly. How do you define \"slightly\" here?\n3. Why are you calling Eqn. (10) a \"complexity measure\"? Which function class is this complexity associated with? To me, Eqn. (10) is just the sum of coding rates at all layers. \n4. Could you also use models that are not CRATE in the experiment of Section 5, and see if SRR can still predict their generalization?\n\n**Summary:** I really like the analysis in Section 4.1 and think that this is a great contribution. However, I also have some big concerns with the subsequent sections. Currently I am rating this paper \"borderline accept\". I encourage the authors to revise this paper during rebuttal. Based on the final version, I could upgrade my rating to \"accept\".", "limitations": "Limitations are not discussed. I encourage the authors to add a limitation section in the revision.\n\n**Post-rebuttal note:** After discussing with my fellow reviewers and the AC, I raised my rating to 7.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "This paper has clear strengths and clear weaknesses. Overall I enjoyed reading this paper, because\n- It is very well written and easy to read. I hardly had any difficulty in understanding this paper.\n- The analysis (especially in Section 4.1) is very clear. It clearly shows the problem of CRATE.\n- The results are new as far as I know.\n- The subject matter studied, which is a white-box deep learning architecture, is very important. I believe that many people will appreciate the analysis in this paper.", "weaknesses": "The paper has some clear weaknesses. I encourage the authors to revise this paper during the rebuttal period, and I'd be glad to provide my comments and feedback during revision.\n\n- What is the point of proposing CRATE-N and CRATE-T? What problems do they solve? From the paper, I cannot understand why these two variants could be better than CRATE. Figures 2 and 3 don't show their superiority over original CRATE. Moreover, I don't see what interesting conclusions could be drawn from the comparison between these variants and the original CRATE.\n- While the authors argue that CRATE can minimize the SRR objective, it seems to me from Figures 2 and 3 that this is not the case. The SRR at the last layer is still close to that at the first layer, especially for the two variants. Of course, SRR does decrease in intermediate layers, but the authors are not suggesting pruning the last layers of CRATE and using the intermediate representations.\n- The experiment in Section 5 and Table 1 is confusing, and I am concerned whether it can be safely concluded from this experiment that SRR is a good proxy of generalization. First, I don't see why this is a \"complexity measure\". In statistical learning theory, a \"complexity measure\" is typically correlated with the size of a certain function class, but I don't see any function class here. Second, all models used in this experiment are CRATE models as shown in Table 3, which presumably optimize SRR. I am worried that the high correlation between SRR and generalization is a special property of CRATE models. To argue that SRR is a good proxy of generalization for all models, the authors ought to use models other than CRATE, that is models that have nothing to do with SRR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1: What is the point of proposing CRATE-N and CRATE-T? What problems do they solve?** \n\nThe goal of developing CRATE-N is to address the potential issue of CRATE-C. As we pointed out that the update (7) in CRATE-C, which performs a gradient descent only for the second-order term, could maximizing $R^c$, rather than minimizing it. This contradicts the principle of SRR. Therefore, CREATE-N is designed by seeking the opposite direction of CRATE-C, which could perform decreases of $R^c$ more faithfully, thus aligning better with the SRR principle. \n\nThe goal of developing CRATE-T is to address the misalignment between CRATE and CRATE-C (i.e., the misalignment between CRATE and the SRR principle). As mentioned in line 107, CRATE replaces the output matrix $U=[U_1,\\dots,U_K]$ in the MSSA operator with learnable $W$ (which is different from $U$). We then raise the following question on the manipulation of the output matrix: if we are free to adjust the output matrix while sacrificing interpretability, can we find more alternatives that can outperform CRATE-C or even CRATE? Therefore, CRATE-T is a feasible choice without introducing new parameters, which can be utilized to better understand the SRR principle and its connection to the performance. We will revise our Section 4 accordingly to improve clarity.\n\n**W1: Conclusions drawn from the comparison with CRATE**\n\nWe want to clarify that our analysis intends to compare the variants with CRATE-C, not CRATE, because CRATE introduces learnable parameters $W$ that are less interpretable. We believe there are at least some interesting conclusions from the comparison: 1) CRATE-N achieves better performance by following the SRR principle more faithfully, shedding light on the connection of SRR to generalization; 2) We need to explore more design choices (e.g., CRATE-T, which may deviate from directly optimizing the SRR but still exhibit a similar architecture) to gain a complete understanding of the SRR principle for model performance (this motivates our Section 5).\n\n**W2: SRR at the last layer close to that at the first layer, especially for the two variants. Authors not suggesting pruning the last layers of CRATE and using the intermediate representations.**\n\nThe SRR objective contains minimizing two terms: $R^c(Z)$ and $\\lambda\\|Z\\|_0-R(Z)$. Based to analysis in Section 4, the update in (8) can decrease $R^c$. Yet, the operation in (3), i.e., ISTA operation, designed to optimize the latter term also matters for the whole objective. This structure may have effects on the rise of curves in Figure 2 and 3. In fact, the sparsity term $\\|Z\\|_0$ has been discovered to shoot up in the last layer in Figure 3 of the original CRATE paper [1]. We also make similar discoveries in other variants (see **Figure 3 in the pdf**).\n\nTo see if this rise affects the intermediate presentations, we present the result of linear probing on them in **Figure 1 in the pdf**. It shows that the representations becomes more linearly separable as the layer goes deeper, although SRR is not well-optimized in the last few layers. This suggests there might be a trade-off between SRR, presumably sparsity, and representation learning in the last few layers. This also means SRR does not necessarily indicate linear separability, especially when the architecture does not faithfully implement the SRR objective. \n\n**W3 & Q3: Why is SRR a \"complexity measure\". Don't see any function class here.**\n\nHere we use \"complexity measure\" as we think higher SRR measure implies more complicated models, then it can be viewed as a \"complexity measure\" for a model. We agree that the rigorous definition of it should be for a certain function class. We will replace it with \"proxy/measure of generalization\" in the revision.\n\n**W3 & Q4: Worried that the high correlation between SRR and generalization is a special property of CRATE models. Could you also use models that are not CRATE in the experiment of Section 5, and see if SRR can still predict their generalization?**\n\nThis is a good question. This is also one of the motivations to consider CRATE-T, which is not intentionally designed to minimize or maximize SRR. This variant can expand the candidate classes of transformer-like models in the correlation analysis. We will use more models such as those replacing the output matrix with randomly initialized and fixed weights or just identity matrix. We will also include more general models, such as standard transformers and CNNs for the experiment in Section 5. We have performed the preliminary analysis for transformers (see **Figure 5 in the pdf**). However, since the correlation analysis requires training a large number of models, we are short of time and computing resources to complete during the rebuttal phase. We will definitely include the experiments and discussions in the revision.\n\n**Q1: Cannot see why CRATE-N makes sense. If it is gradient ascent, how does it minimize $R^c$?**\n\nCRATE-N actually performs gradient ascent on the second-order term in (5). This corresponds to the green curve in Figure 1(b), which tends to make the eigenvalue $\\lambda$ smaller (as the optimization is performed for $\\lambda\\ge1$ and small $\\lambda$ increases the green curve). On the other hand, the right-hand side of (5) itself, which can be understood as the orange curve in Figure 1(b), has an opposite pattern as the green curve, which will decrease when $\\lambda$ becomes smaller enough. This leads to the minimization of $R^c$ (which can be expressed with a logarithmic function of eigenvalue $\\lambda$).\n\n**Q2: In line 170 the authors claimed that SRR \"...rises slightly...\". But it seems to me that it rises significantly.**\n\nWe will make our description more accurate in the revision.\n\n**Add a limitation section in the revision.** We will add a limitation section in the appendix in the revision. \n\n[1] Yu et al. White-Box Transformers via Sparse Rate Reduction. 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": "OCjEWCoRvg", "reviewer": "Reviewer_uGk9", "summary": "This paper considers a recent line of transformer-like models called CRATE where each layer is designed to approximate a gradient-based optimization step of an information-theoretic objective function called sparse rate reduction. The contributions of the paper are: (1) investigating whether CRATE actually implements optimization of sparse rate reduction, and proposing variants of CRATE based on the observations, and (2) empirically analyzing the relation between the measured sparse rate reduction in CRATE (and the proposed variants) and their capabilities for generalization. The paper first shows that the approximation involved in the original derivation of CRATE results in cases where the forward pass of each layer does not implement gradient descent on sparse rate reduction, but does the opposite (ascent); from this observation, the authors propose CRATE-N and CRATE-T, where the former negates the output of a layer (before applying skip connection) and the latter transposes the output projection. Then, on CRATE and the proposed variants, the authors study how sparse rate reduction actually happens, and proceeds to studying the (causal) relationship between sparse rate reduction and generalization by measuring their correlation followed by showing that explicitly regularizing sparse rate reduction on the final layer leads to improved generalization.", "questions": "- Q1. In Equation (12), how is the generalization gap $g(\\theta)$ measured? Is it measured as the difference between the training loss (0.01) and test loss?", "limitations": "The authors have partially discussed the limitations in the checklist. I encourage the authors to create a separate section in Appendix.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- S1. The paper studies an interesting question of whether sparse rate reduction correlates with generalization in CRATE, among other questions studied in the paper. The proposed variants (CRATE-N and CRATE-T) are technically original as far as I am aware.\n- S2. The results given in Table 1 that sparse rate reduction correlated well with generalization gap for models is interesting, in particular since it outperforms a sharpness-based measure which is known to correlate with generalization.\n- S3. The paper is overall well written and easy to follow.", "weaknesses": "- W1. The motivations and technical soundness of the proposed variants of CRATE is unclear, especially given that they underperform on CIFAR-10/100 compared to the original CRATE as shown in Table 2. For CRATE-C, the authors point to Equation (7) and differentiate it from the original CRATE, but from the main text it seems Equations (5-7) describe the approximations involved in the original CRATE and how CRATE-C is different is not clear. Also, in Line 125, the authors state that \"update step takes the gradient of a lower bound of $R^c(\\mathbf{Z};\\mathbf{U})$\" which was a bit weird to me considering that the objective is minimizing $R^c(\\mathbf{Z};\\mathbf{U})$; I am not aware of optimization approaches that *minimize the lower bound* of the objective function, although this might be due to my limited knowledge in optimization literature. For CRATE-N, the authors claim in Lines 147-149 that \"taking gradient ascent on the second-order term in the Taylor approximation of $R^c(\\mathbf{Z};\\mathbf{U})$ minimizes the eigenvalues and consequently $R^c(\\mathbf{Z};\\mathbf{U})$\", which is not true for general objective functions as far as I am aware - please correct me if I am missing something. For CRATE-T, the authors motivate it in Line 152 by saying that \"replacing the output matrix by learnable parameters is problematic ... it also breaks the inherent structures and sacrifices the mathematical interpretability\", but then proceed to still use learnable output matrix, only transposing it as far as I understand. It is unclear to me what particular arguments the authors are trying to show with this variant.\n- W2. In Section 4.3, the reported tendency of increasing sparse rate reduction measure after certain depth (Figure 2 and 3) does not agree with Figure 4 of the original CRATE paper [1], but proper discussion is not given. Also, although the authors have motivated CRATE-N (among other variants) from the pitfalls of the formulations of the original CRATE, it shows a similar tendency of worsening sparse rate reduction after certain depth. The experiments themselves are conducted on CIFAR-10/100, which is more limited compared to the similar analysis in Section 3.1 of the original CRATE paper [1] which has used ImageNet-1k.\n- W3. In Section 6.2, in Line 271, it is not clear what the authors meant by \"regularizing the training loss with sparse rate reduction for each layer ... lacks fine-grained controls over parameters\". Regarding the results, the performance gain obtained by regularizing the last layer is marginal, and it is unclear if the gain is statistically significant, and whether it outweights the computational cost of computing the measure of sparse rate reduction. In fact, Table 5 shows that regularizing layer 4 achieves the best performance gain overall; the reason the authors have particularly chosen layer 12 to be shown in Table 2 is unclear. Lastly, Table 5 shows that random layer regularization using sparse rate reduction leads to degraded performances compared to not using regularization, and the reason is not clear: I believe underfitting is not likely as the dataset is CIFAR-10, which implies that random layer regularization using sparse rate reduction hurts generalization.\n\n[1] Yu et al. White-Box Transformers via Sparse Rate Reduction (2023)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1:Unclear motivations/soundness of variants**\n\nCRATE-N aims to counteract issues in CRATE-C where the update can increase $R^c$, opposing SRR principle. By moving in the opposite way of CRATE-C, CRATE-N implements decrease in $R^c$ more faithfully, aligning better with SRR principle. CRATE-T addresses the discrepancy between CRATE and CRATE-C by questioning the replacement of the output matrix in MSSA operator with learnable parameters in CRATE. By exploring alternatives without new parameters, CRATE-T aims to enhance comprehension of SRR principle and improve performance of CRATE-C.\n\n**W1:Authors point to (7) and differentiate it from CRATE...how CRATE-C is different?**\n\nSorry for making you confused. (5-7) describe the approximation and derivation of part of CRATE-C, and (7) is the update in CRATE-C. On the other hand, CRATE involves further modification-replacing matrix $U=[U_1,\\dots,U_K]$ with learnable parameters $W$(line 106). The title of Section 4.1 should refers to \"CRATE-C\", not \"CRATE\".\n\n**W1:\"...update step takes the gradient of a lower bound...\" weird as the objective is minimizing $R^c$**\n\nYou are right, minimizing the lower bound of the objective cannot reliably lead to the minimization of the original objective $R^c$. This is exactly the issue we want to highlight for CRATE-C. \n\n**W1:\"taking gradient ascent...consequently $R^c$\", not true for general objective functions**\n\nFor the general function, you are correct. Yet, we're not using a general function. As in line 122, eigenvalues of a PSD matrix plus an identity matrix are definitely no less than one. When $\\lambda \\ge 1$, gradient ascent on the second-order term of $R^c$ (green curve in Figure 1(b)) will decrease $\\lambda$, hence decreasing $R^c$ (blue curve). Therefore, our statement still holds.\n\n**W1:Authors motivate CRATE-T in \"replacing output matrix...breaks the structures...\", but then still use learnable output matrix, unclear arguments to show**\n\nThe reason for sacrificing the interpretability is not about \"learnable\", it is about learnable parameters are \"different\" for output matrix and Q/K/V matrices. This can greatly enhance the performance. We then want to explore other designs not entirely derived from theory. For example, we have experimented with fixed weights and identity matrix in **Table 1 in the pdf**. We want to highlight designs different from $[U_1,\\dots,U_K]$ in CRATE-C without new parameters.\n\n**W2:Increasing SRR after certain depth (Figure2/3), not agree with Figure4 of [1], need discussion.**\n\nThere are several reasons: 1) Most importantly, different y axis. They plot the figures of compression and sparsity term separately, which fail to reflect the true tendency of compression term because they use $Y^\\ell$ in (2). Instead, we visualize SRR measure as a whole and use $Z^\\ell$ in (3); 2) They apply $\\ell_2$ normalization before measurement of $R^c$ while we follow the definition. 3) CRATE-Base v.s. CRATE-Tiny configurations; 4) ImageNet-1k v.s. CIFAR10/100. \n\n**W2:Motivated from pitfalls of the original CRATE, CRATE-N still shows worsening SRR after certain depth.**\n\nSRR has two terms: $R^c(Z)$ and $\\lambda \\|Z\\|_0-R(Z)$. Based on analysis in Section 4, the update in (8) can decrease $R^c$. Yet, the operation in (3) that aims to optimize the latter term also matters for the whole objective and may affect $R^c$ as well. This structure could have influences on the rise of the curves in Figure 2 and 3. How the ISTA block in (3) interacts with the MSSA block is an open question.\n\n**W2:Limited experiments compared to [1] that used ImageNet-1k.**\n\nWe provide SRR measure similar to Figure 2 and 3 on ImageNet-1k in **Figure 2 in the pdf**. The models are trained with the official recipe [1]. We also provide the sparsity term in SRR measure $\\|\\boldsymbol{Z}\\|_0/(d*N)$ in **Figure 3 in the pdf**. The sparsity term has a similar tendency with the SRR measure, and it goes up in the last few layers, which means that ISTA block also matters and need further research.\n\n**W3:Unclear what \"lacks fine-grained controls over parameters\" means.**\n\nWhat here emphasize the reg term in (10) is an average SRR measure over all layers, and is highly inefficient to compute. The SRR measure at some layer might be more important and should be optimized in isolation. Therefore we propose to regularize one specific layer at each iteration. We'll make this point clear in the revision.\n\n**W3:Reg on the last layer, marginal performance gain.**\n\nOur goal is NOT to demonstrate its superiority in performance gains. Instead, we want to complement the conclusions in Section 5. SRR measure is shown to be better than sharpness, then it should be reasonable to incorporate it into the training for generalization, similar to sharpness-aware minimization [2]. One direct approach is through regularization. We only provide preliminary results and leaves room for engineering to improve the results.\n\n**W3:Why layer 12 in Table2.**\n\nSpecifying which layer to regularize could be computationally prohibitive, especially when the model size grows. We intuitively select the last layer, which should be reasonable if the depth of the models scales. Results in Table 2 indicate that this intuitive choice can already give consistent performance gains in different settings.\n\n**W3:Random layer reg in Table5 performs worse than w/o it, reasons unclear.**\n\nWe also provide the accuracy on the training set of CIFAR10 in **Table 2 in the pdf**. The performance drops more on the training set than validation set. So underfitting could be the problem that hurts the generalization.\n\n**Q1:How is the generalization gap measured?**\n\nIt's the difference between validation and training loss at convergence (training loss reaches 0.01).\n\n**Limitations.**  We will include the limitation section in the revision. \n\n[1] Yu et al. White-Box Transformers via Sparse Rate Reduction. \n[2] Foret, et al. Sharpness-aware Minimization for Efficiently Improving Generalization.", "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": "UG26zLYWMj", "reviewer": "Reviewer_SqrQ", "summary": "The paper investigates a Transformer-like deep network architecture CRATE based on algorithmic unrolling of the sparse rare reduction objective function. It points out some pitfalls of the approximated layer operation where it does not decrease the R^c as it should, studies some alternatives, and also points out some surprising phenomena of training the unrolled architecture. Finally, it investiages whether SRR can be used as a regularizer in the loss function to promote generalization performance.", "questions": "See above", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The investigation pointed out many interesting observations. For example,\n\n\n\n1. the softmax attention update can be seen as an approximation to the lower-bound of R^c, an alternative view to the CRATE paper\n2. Without Relu, the softmax attention update increases R^c as opposed to decreasing R^c, while with Relu added R^c decreases. \n3. As the training proceeds, the SRR measure at any layer decrases as opposed to increases. \n4. The SRR measure has the strongest correlation of the generalization gap (albeit i do not understand its definition, as illustrated in weakness), which motivated adding it to the learning objective as a regularizer for generalization.", "weaknesses": "Despite the interesting observations, there seem to be disconnections here and there. \n\n## Section 4.1:\n\n\n\n1. Lines 125 and 135-136 make conclusions based on how (6) is related to (7), but these conclusions are not justified. \n\n    _It can be seen that this update step takes the gradient of a lower bound of R c (Z; U) and discards the first-order term._\n\n* It is true that from (6) to (7) the first order term is removed as the authors claimed. Note however a softmax is also added. \n\n    _When taking the gradient of R c to construct the MSSA operator, omitting its first-order term will produce a counterproductive effect._\n\n* Maybe omitting the first-order term is not the (only) culprit, adding softmax could also play a role.\n2. Line 141: _Therefore, one step of update (7) secretly maximizes R c_. I’m not sure about this either. Note that the green and orange curves are both quadratic functions, and are shifted/scaled versions of each other. Why would doing gradient descent with the orange be better than with the green?\n\n    I would imagine at least two potential answers. 1) It depends on where gradient descent is initialized, i.e., at the left or right to the peak of the curve. But this needs to be justified better. 2) Using a quadratic approximation of the log of eigenvalues as in (6) is already problematic.\n\n3. I think part of the issues above can be addressed by making a more complete experiment. More precisely, ​​using the same parameters settings as in making Figure 1(a), we should have curves resulting from\n    1. the original gradient update\n    2. the update from 2nd order Taylor expansion, i.e., equation (6)\n    3. the update from removing the second order term from (6)\n    4. the update from further adding softmax, i.e., (7)\n\n## Section 4.2:\n\n\n\n1. I do not understand the motivation of the variants proposed in the 4.2. ​​Granted, (7) is not the best approximation to the SRR. How do (8) and (9) make better approximations? It seems that they are just arbitrary modifications from (7) not serving any goals.\n\n    If the argument, as the author states in lines 135-136, is that omitting the first-order terms in (6) is bad, then isn’t it natural to see what if one keeps the first-order terms?\n\n\n## Section 5.6\n\n\n\n1. Where is the generalization gap g(theta) defined? I do not seem to find them. I presume mu is any complexity measure, and tau is Kendall’s coefficient that is cited.\n2. In table 2, the change of the accuracies by adding SRR regularization seems relatively small. I do not know if there is significance in it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1: Lines 125 and 135-136 make conclusions based on how (6) is related to (7), but these conclusions are not justified.**\n\nThe derivations in (6-7) are rephrased from the original CRATE paper [1]. The softmax introduced after omitting the first-order term is quite intuitive, which converts auto-correlation to the distribution of membership of different subspaces. We refer the reviewer to Section 2.3 of [1] for discussion.\n\n**W1: Note however a softmax is also added from (6) to (7). Maybe omitting the first-order term is not the (only) culprit, adding softmax could also play a role.**\n\nYour concern that introducing softmax may also affect the optimization is understandable. But note that, as we mention in line 128, the update we use in the toy experiment is (7), which contains the softmax. Therefore, Figure 1(a) should be able to corroborate our analysis, and adding softmax does not change the conclusion.\n\n**W1: Therefore, one step of update (7) secretly maximizes $R^c$. I’m not sure about this either. Note that the green and orange curves are both quadratic functions, and are shifted/scaled versions of each other. Why would doing gradient descent with the orange be better than with the green?**\n\nThe two answers you mention are reasonable. But we are not arguing which one is better. After all, the goal is to minimize the blue curve$-$the log function. If using the orange curve, we need to constrain the eigenvalues on the left of the peak, which is hard. The update (7) discards the first-order term, corresponding to the green curve. Now, minimizing the green curve makes the eigenvalues $(\\ge 1)$ increase, hence maximizing the blue curve which corresponds to $R^c$. The quadratic function is the most straightforward approximation for designing the theory-grounded structure$-$MSSA operator. By optimizing other approximations of log functions, more powerful and more mathematically meaningful structures may emerge.\n\n**W1: part of the issues above can be addressed by making a more complete experiment**\n\nWe follow these settings and provide the visualization in **Figure 4 in the pdf.** The results shows that the original gradient update of $R^c$ (a) with no approximation will monotonically decrease $R^c$, which is expected. Using the Taylor expansion approximation (b) will increase $R^c$, meaning that the eigenvalues $\\gamma$ are mostly greater than one and the optimization happens to the right of the peak of the orange curve in Figure 1(b). Removing the first-order term of the Taylor expansion (d) monotonically increases $R^c$ which is the issue we specify in CRATE-C. Finally, further adding the softmax (e), leading to the update in (7), will not change the monotonic increase of $R^c$. We believe this figure will address most of the concerns in this part. \n\n**W2: I do not understand the motivation of the variants proposed in the 4.2. ​​Granted, (7) is not the best approximation to the SRR. How do (8) and (9) make better approximations?**\n\nThese two variants have different motivations. CRATE-N is motivated to counteract the issue in CRATE-C where the update (7) increases the coding rate $R^c$, conflicting with the SRR principle. By moving in the opposite direction of CRATE-C, CRATE-N with (8) aims to implement the decrease in $R^c$ more faithfully, aligning better with the SRR principle. On the other hand, CRATE-T is not motivated for better alignment with the SRR principle. Instead, it is motivated by designing competitive alternatives of CRATE-C or even CRATE without introducing new parameters, since CRATE replaces the output matrix $U=[U_1,\\dots,U_K]$ in CRATE-C with learnable $W$ (line 107) which sacrifices interpretability. \n\n**W2: If the argument, as the author states in lines 135-136, is that omitting the first-order terms in (6) is bad, then isn’t it natural to see what if one keeps the first-order terms?**\n\nOmitting the first-order term is bad when doing gradient *descent* because this will maximize $R^c$, but it is good when doing gradient *ascent* since this minimizes $R^c$ which is (8). It’s indeed natural to keep it, but we have to constrain the eigenvalues on the left to the peak of the orange curve, which is difficult. \n\n**W3: Where is the generalization gap $g(\\theta)$ defined?** \n\nWe measure the generalization gap as the difference between validation loss and training loss at convergence (training loss reaches 0.01), i.e., $L_{val}(w)- L_{train}(w)$. We can give a formal definition in the appendix in the revision. \n\n**W3: In table 2, the change of the accuracies by adding SRR regularization seems relatively small.**\n\nOur goal is not to demonstrate the superiority of SRR over other kinds of regularization or for substantial performance gains. Instead, we would like to validate and complement the conclusions in Section 5. SRR measure is shown to be better than sharpness as a predictor of generalization, then it should be reasonable to incorporate it into the training for improved generalization, similar to sharpness-aware minimization [2]. One direct approach is through regularization. We only provide preliminary results as a proof-of-concept. We believe there is room for engineering to make the results more significant.\n\n[1] Yu et al. White-Box Transformers via Sparse Rate Reduction. NeurIPS, 2023.\n\n[2] Foret, et al. Sharpness-aware Minimization for Efficiently Improving Generalization. ICLR, 2021.", "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": "z6i8b4DmpT", "reviewer": "Reviewer_7qne", "summary": "The work investigates the design choices and analysis components in the recently proposed white-box transformer (CRATE). The work studies and analyzes the design rolling objective (SRR) in CRATE-like architecture.  The work also proposes using SRR as layerwise constraints and shows the SRR will further improve the model performance especially on the variants proposed in the paper. The work also investigates using SRR as a measure for generalization.", "questions": "NA", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The work is well-written, with clear structures and language.\n2. The work provides thorough studies into the design choices and justification of SRR as unrolling objectives. \n3. The findings in the paper provide valuable additional insights to the original CRATE paper. \n4. The analysis on generalization via Correlation of SSR is valuable that could potentially encourage future works.", "weaknesses": "1. The analysis in section 5 is very interesting. It will benefit the paper if the authors can provide more insights and investigation. For example, examining using SSR to measure other pretrained transformer-like architectures. \n2. The analysis in Section 4 and experiments in Section 6 uses different alpha values. It could further improve the work if consistent alpha is used.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**It will benefit the paper if the authors can provide more insights and investigation from section 5. For example, examining using SSR to measure other pretrained transformer-like architectures.**\n\nEvaluating $R^c(Z;U)=\\sum_{k=1}^K \\frac{1}{2}\\log \\operatorname{det}(I+\\gamma (U_k^TZ)^T (U_k^TZ))$ in SRR need to specify $U$. However, it is not well defined on models other than CRATE and its variants like general transformers. Moreover, conducting correlation analysis needs to train a large number of models, but currently we lack the resources and time to accomplish it. Despite that, we will provide some preliminary results on evaluating SRR on a standard transformer in **Figure 5 in the pdf**. Specifically, we set $U=[U_1,\\dots,U_K]$ equal to the query, key, and value projections respectively. This shows we may generalize the usage to a standard transformer, or even more general models. We will surely add this experiment and discussion in the revision. \n\n**The analysis in Section 4 and experiments in Section 6 uses different alpha values. It could further improve the work if consistent alpha is used.**\n\nWe will change the Figure 1(a) using consistent $\\alpha$ with that in Section 6 in the revision.", "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:
We sincerely appreciate the thoughtful reviews and comments provided by all reviewers. Below, we address the main points raised, details can be found in corresponding blocks for each reviewer:

- Reviewer RBPW questioned the role of different variants and the behavior of the SRR objective. We clarify that the variants have different motivations and highlight the benefits of adhering closely to the SRR objective and its relation to intermediate representations in the linear probing experiment.

- Reviewer uGk9 primarily raised confusion and concerns on the specification of problems and called for the use of complex datasets. We address the concerns regarding the performance in the experiment and elucidate the potential reasons for particular phenomena with the help of a larger dataset. Clearer descriptions are given to explain the details on the analysis.
  
- Reviewer SqrQ suspected the effects of approximation on the design goal and pointed out the missing definition and discussions. We provide similar toy experiments with different settings to illuminate the problems and make clear the interpretation regarding the figures.

- Reviewer 7qne acknowledged the significance of the experimental results and provided suggestions to improve the work.

Meta-review: The authors conduct an in-depth investigation of CRATE, a previously proposed Transformer-like architecture derived from unrolled optimization of an information-theoretic objective function defined on representations referred to as “sparse rate reduction” (SRR), with subsequent approximation. The authors present a careful analysis of CRATE, showing in a toy experiment that these approximations (made for computational efficiency and performance reasons) actually lead the network to increase the SRR from layer-to-layer, rather than decrease it as per the design, and proposing architectural variants (“CRATE-N” and “CRATE-T”) designed to correct this discrepancy without sacrificing downstream performance. To better understand the remaining gap between CRATE’s performance and the performance of the proposed variants, they propose to use SRR as a complexity measure for predicting generalization; experiments on several CRATE model variants and transformers (post-rebuttal) highlight that SRR correlates the best among different generalization measures (e.g., sharpness) with the generalization performance.

After the rebuttal, in which authors clarified many vague points and presented new experiments corroborating their arguments on the relationship between SRR and generalization on non-CRATE architectures, scores for the paper increased, converging on accept. Reviewers appreciate the careful design of the paper’s experimental and theoretical studies and the additional insight it yields towards understanding generalization in transformer-like architectures. The AC concurs with the reviewers’ judgment, and recommends acceptance. The authors are encouraged to incorporate all feedback from the reviewers into the revision, especially regarding additional promised experiments on non-CRATE variants.
Decision: Accept (poster)