An In-depth Investigation of Sparse Rate Reduction in Transformer-like Models

Authors: yunzhe  hu
difan  zou
dong  xu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:48:34

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.

Main paper: 
 introduction: 
 Transformers [39, 11] have become the de facto choice of 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 a mathematical perspective. Works in this vein have attempted to establish connections between Transformers and a reverse-engineered energy function [43, 17] , associative memory such as modern Hopfield network [35, 37, 26, 4] and sparse distributed memory [5] , or programming languages [41, 22] , to 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), a function that promotes maximum information gain described by the coding rate function [24, 46] and induces sparsity. In particular, they show that a Multi-head Subspace Self-Attention (MSSA) operator with skip connection and an Iterative Shrinkage-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 a Transformer-like model CRATE in which every layer should have the completely interpretable compress-then-sparsify behavior. However, although motivated by an information-theoretic 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 as 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 a counterproductive 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 as a complexity measure has a positive correlation coefficient and outperforms other baselines.• In Section 6, we demonstrate the effectiveness of SRR as a regularization technique for improved performance on 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 
interpreting transformers: 
 Research on interpreting Transformers [39, 11] 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 [14] , visualizing attention map [7, 8, 44] , etc. Several works opt for "mechanistic interpretability" [12, 28, 40] aiming to reverseengineer 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 as Bayesian inference [1] , convex optimization [36] to analyze attention in Transformers. There have also been attempts to interpret a Transformer as an energy function optimizer [43, 17] , connect attention to memory [35, 37, 26, 4, 5] or interacting particle systems [13] or transform into humanreadable programs [41, 22] , to name just a few. Our work focuses on the empirical investigation of a Transformer-like model, CRATE [45] , recently introduced from pure mathematical derivation.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 a forward-constructed manner. For instance, the seminal work [15] proposed to unroll the Iterative Shrinkage-Thresholding Algorithm for sparse coding into layers of linear operation followed by ReLU non-linearity. Other works have tried to find a 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 ReductionLet Z = [z 1 , . . . , z N ] ∈ R d×N denote N samples, where each column z i ∈ R d represents tokens in Transformers. U = [U 1 , . . . , U K ] ∈ R d×Kp denote a set of incoherent basis spanning K subspaces, wherein columns of U i ∈ R d×p represent basis in i-th low-dimensional subspace (p < d).We follow the configuration that d = Kp as in standard ViT [11] .Previously, Yu et al. [46] propose that the compactness of representations Z ∈ R d×N can be measured by a coding rate function: R(Z) .= 1 2 log det(I + d N ϵ 2 Z T Z).A more recent study [45] contends that the objective of representation learning is to transform and compress samples from an unknown distribution to a mixture of low-dimensional Gaussian distributions supported on incoherent bases. This objective boils down to the maximization of Sparse Rate Reduction (SRR): maxEQUATIONwhere∥ • ∥ 0 means ℓ 0 norm and R c (Z; U ) . = K k=1 R(U T k Z)measures the compactness of representations in the low-dimensional subspaces. One layer of a network, formulated as a mapping f w (•) parameterized by w, can be interpreted as applying one step of gradient-based methods to the objective in (1) . In practice, Yu et al. [45] use alternating minimization to break down the optimization into two steps: compression, i.e. min Z R c (Z; U ) and Sparsification, i.e. min Z λ∥Z∥ 0 -R(Z). Specifically, given representation Z ℓ-1 at (ℓ -1)-th layer, Z ℓ can be obtained by two-step optimization:EQUATIONEQUATIONwhere α, β > 0 are step sizes, D ℓ ∈ R d×d is assumed as a complete dictionary, scalar γ . = p N ϵ 2 andEQUATIONThe operator MSSA(•; U ) in ( 4), called the Multi-head Subspace Self-Attention (MSSA) operator, takes the form of self-attention in standard Transformers [39, 11] , with tied query, key and value matrix, i.e., U T k while the output matrix being its transpose, i.e. U k . Instead of strictly following this formulation, they further replace [U 1 , . . . , U K ] ∈ R d×Kp in the MSSA operator with an additional learnable parameter W ∈ R d×Kp . To distinguish them, we name the model with implementation (4) CRATE-C(onceptual). By incrementally optimizing (1) with alternating minimization, a Transformer-like model with layered structures can be naturally constructed. With input Z 0 , e.g. tokenized images in ViT, an L-layer model iteratively optimizes the input and yields the final representations Z L . Parameters {U ℓ } L ℓ=1 and {D ℓ } L ℓ=1 can be learned through end-to-end training [15] . 
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 crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
is sparse rate reduction optimized in transformer-like models?: 
 While the white-box Transformer-like model proposed in [45] is derived by unrolling optimization upon a pre-defined objective function, whether the optimization is implemented by the model in the forward pass is still unclear. In this section, we first review the main derivations at the core of building CRATE, i.e. unrolling optimization min Z R c (Z; U ) into MSSA operator with skip connection as in (2) , and identify the pitfalls in implementing the minimization. We then provide variant models based on different implementations and empirically show their layer-wise behaviors. Figure 1 : In a simplified attention-only experiment, MSSA operator with skip connection actually implements an ascent method on R c (Z; U ), opposed to its design purpose (left). This is due to an artifact in approximation with its second-order term. (right)We first show that the second-order Taylor expansion of the coding rate of representations Z projected onto subspaces can be expressed as:EQUATIONwhereλ k i ≥ 1, i ∈ [N ] are the eigenvalues of I + γ(U T k Z) T U T k 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 c (Z; U ):EQUATIONEQUATIONEQUATIONIt 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. With a proper step size, the coding rate on the same subspaces is expected to 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 L layers of update (8) are conducted with parameters {U ℓ } L ℓ=1 initialized as orthonormal matrices. We initialize a random variable Z 0 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 a proper ϵ 2 such that γ = 1. As shown in Figure 1a , R c (Z ℓ ; U ℓ ) is always greater than R c (Z ℓ-1 ; U ℓ ) and R c (Z ℓ ; U ℓ ) is increasing in general as the layer goes deeper. This means the update (8) that resembles the standard self-attention with skip connection does not essentially implement a descent method on R c . The crux lies in the approximation of R c 's gradient.When taking the gradient of R c to construct the MSSA operator, omitting its first-order term will produce a counterproductive effect. As shown on the left-hand side of the inequality in ( 5), R c can be expressed as the sum of logarithms of eigenvalues. We expect the eigenvalues to decrease to minimize the value of R c . Figure 1b illustrates 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 -(λ k i -1) 2 2), the eigenvalues will go up leading to an increase in the value of R c . Therefore, one step of update (8) secretly maximizes R c , contrary to the purpose of its design. More figures detailing this issue are in Appendix A.In the previous subsection, we show the problems arising from gradient approximation when unrolling min Z R c (Z; U ) into MSSA operator with shortcut. 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-induced architectures but in a more self-contained way. They also serve as 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 R c , therefore minimizing the eigenvalues and consequently R c . This variant is designed to counter the pitfalls in CRATE-C, enabling a more faithful reduction in R c and thereby enhancing alignment with the SRR principle. We term the Transformer-like model with this implementation CRATE-N(egative):EQUATIONThe other variant we would like to introduce is motivated by the misalignment between CRATE and CRATE-C. Although replacing the output matrix [U 1 , . . . , U K ] 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? It turns out that a 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(ranspose):EQUATIONThe Transformer-like model CRATE is built by sequentially stacking the layer that comprises two modules in ( 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 is 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λ∥Z ℓ ∥ 0 + R c (Z ℓ ; U ℓ ) -R(Z ℓ) at different layers and epochs on the training set. λ is chosen as 0.1 and detailed experiment settings can be found in Section 6.1. 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 4.1 we demonstrate that R c (Z; U ) will monotonically go up in the absence of operation (3) . We conjecture that the ReLU non-linearity may also play an important role in optimizing the compression term R c (Z; U ) in the forward pass. Another surprising finding is that as the learning process proceeds, the sparse rate reduction measure at each layer will increase monotonically across all models, with a rare exception in the last few layers of CRATE-C.These phenomena give us implications for understanding Transformer-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 in a 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-inspired operator MSSA diverges from its goal from a micro perspective. We postulate that ReLU non-linearity in (3) could also promote compression and leave their interaction for future work.5 Whether Sparse Rate Reduction Benefits Generalization?So far, we have partially confirmed the validity of different implementations of Transformer-like models by inspecting the layer-wise optimization of SRR. But whether this objective is 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 to the generalization of CRATE.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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. 
sparse rate reduction as a 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 a smaller generalization gap. Complexity measures can be either theoretically motivated, such as PAC-Bayes [25, 29] , VC-dimension [38] , norm-based bounds [32, 3, 30] or empirically motivated, such as sharpness [20] and path-norm [31] . We choose to adapt SRR into a complexity measure that belongs to the latter category:EQUATIONwhere Z ℓ denotes the output at layer ℓ and w ℓ contains the parameters including U ℓ and D ℓ .An effective measure of complexity should bound the generalization gap, defined as the difference between validation loss and training loss when the latter reaches a threshold, i.e., L val -L train , with high probability. However, for those measures that do not provably bound this gap, as is the case with SRR measure (11) , we need to evaluate its correlation with the generalization gap to understand its causal relationship to generalization.Collecting Trained Models To evaluate the complexity measure and generalization across models, we consider changing the hyperparameters and collect a set of models trained to meet a specific stopping criterion. Here, we also consider the model type containing different variants of CRATE as a hyperparameter to investigate its influence on generalization. Formally, let Θ i denote a type of hyperparameter with |Θ i | different choices, and define θ .= (θ 1 , θ 2 , . . . , θ n ) ∈ Θ 1 × • • • × Θ nas an instantiation from n types of hyperparameters. By varying choices across hyperparameter space, we can produce|Θ 1 | × • • • × |Θ n | 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 on CIFAR-10 dataset, when cross-entropy loss reaches 0.01 following the stopping criterion in [18] . Experimental details and choices of hyperparameters can be found in Appendix C.Evaluation Criterion A common method for measuring correlation is 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 a given complexity EQUATIONEQUATION)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, we present the results in terms of Kendall's coefficient τ in Table 1 and scatter plot of SRR measure in Figure 4 when the width d is chosen as 384. The results when d = 768 is deferred to Appendix D. The granulated coefficient Ψ is also reported (see [18] for a detailed definition). We confirm the findings from prior works that some norm-based measures, such as sum/prod of spectral/Frobenius norm of parameters negatively correlate with generalization, even on Transformer-like models. An interesting finding is that path-norm also negatively correlates with generalization, which partly contradicts the previous conclusion. This implies that regularization on path-norm, e.g. Path-SGD [31] , may not be applicable for improved generalization on Transformer-like models. Among the measures we investigated, the inverse of margin and sharpness-based PAC-Bayes flatness show positive and strong correlations. This result justifies the common belief that larger margin or flatter loss landscape leads to better generalization across the investigated Transformer-like models. Compared to baselines, the SRR measure in (11) achieves the highest overall coefficient and, particularly in the model type axis, outperforms the rest. This motivates the use of SRR as regularization in the loss function to improve generalization. 
sparse rate reduction as regularization: 
 Since SRR measure enjoys a strong correlation to the generalization of Transformer-like models, we would like to investigate its potential as the direct regularization to the standard training loss. In particular, we add the SRR measure in (11) by a regularization factor η to the cross-entropy loss:EQUATIONwhere λ > 0 is the regularization coefficient andZ ℓ StopGrad = f w ℓ (StopGrad(Z ℓ-1)). The operator StopGrad here, implemented as "Tensor. detach()" in PyTorch, prevents gradient propagation from the output Z ℓ to the previous layers. This allows parameters w ℓ at each layer to be updated without interfering with each other, giving more precise optimization of SRR in separate layers.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 α = 1, and scaling factor γ = 1. We also include LayerNorm before each operation in (2)(3) for better trainability and learnable positional encoding. A trainable [CLS] token is prepended to the representations for computing cross-entropy loss and classification.We use CIFAR-10 and CIFAR-100 datasets for training and evaluation. In practice, we adopt Adam [21] optimizer and initialize learning rate as 1 × 10 -4 with cosine decay. All models are trained for 200 epochs with batch size as 128. Note that we only use the basic data augmentations: random resize and cropping, horizontal flipping, and RandAugment [9] (with the number transformations n = 2 and magnitude m = 14). We do not use other techniques for state-of-the-art performance but to demonstrate the effectiveness of SRR as regularization. We tune the factor η via a grid search over {0.0001, 0.001, 0.01, 0.1, 1} and find that 0.001 works best. All experiments are conducted on NVIDIA GeForce RTX 3090. Regularizing the training loss with sparse rate reduction measure (11) needs to compute R(Z) and R c (Z; U ) for every layer. However, this is highly inefficient as it involves high-dimensional 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 a pre-defined layer or a random layer with uniform probability during training. In practice, we find that the former works better. Table 2 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 E 
conclusion: 
 To further research in interpreting neural architecture, we provide an in-depth investigation of a recent mathematically driven Transformer-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-like models. Furthermore, we demonstrate its positive correlation to generalization and effectiveness over baselines.Driven by this connection, we show a simple way to use SRR as regularization to improve performance on CIFAR-10/100 datasets. Future direction may include applying layer-wise training and connecting SRR to the Forward-Forward algorithm [16] , or exploring the impact of depth in the unrolled models. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This is no such risk. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes]  
 
 Appendix: This study has several limitations. Firstly, the conclusion that the SRR measure can be a strong indicator of generalization is limited to the CRATE family. Generalizing this conclusion to standard Transformers would be non-trivial, as the SRR measure is not properly defined when the querykey-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 a simple manipulation (CRATE-T). This calls for a rigorous inspection of each component's functionality in the framework. Lastly, while we confirm the positive correlation to generalization, our analysis is limited in scale. Consequently, drawing definitive conclusions regarding whether SRR can be a principle or necessitates further engineering to push the model's limit is challenging. A Better and more systematic way is needed to determine whether SRR is principled for designing the Transformerlike models and quantify this relationship in an appropriate task, perhaps beyond classification.To give a clearer picture of how approximations affect the optimization of R c , we provide the complete results with different update rules under the same simplified settings in the main text: The results in Figure 5 correspond to the above experiments. Gradient descent on R c did make it decrease across layers. Conversely, applying gradient descent on its second-order Taylor expansion resulted in an increase, indicating a potentially flawed approximation. Isolating gradient descent to the second-order term led to a rise in R c , as opposed to the design purpose. Furthermore, incorporating the softmax function, a real-world operation examined in the main text, did not alter this conclusion. EQUATIONEQUATIONFigure 5 : (a) Original gradient update, i.e, (6) (b) update from second-order Taylor expansion, i.e., (7) (c) update from removing the second-order term from (7) (d) update from removing the first-order term from ( 7) (e) update from further adding softmax, i.e., (As mentioned in Section 3, CRATE replaces the output matrix U = [U 1 , . . . , 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? 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.We 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 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).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.It is therefore crucial to isolate the influence of data augmentations from the change of other hyperparameters. We choose to remove data augmentations during training to ensure that most models can be trained to 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 5 sets of hyperparameters, i.e., batch size, initial learning rate, width, dropout probability, and model type. We present the choices of these hyperparameters in Table 4 . Adam [21] 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 in MSSA operator. Table 5 and Figure 6 give results on correlation to generalization when width d = 768. We see that SRR is slightly better than other baseline measures in terms of overall τ . In the axes of dropout and model type, however, it underperforms 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.Table 6 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 to 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.• 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example 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, we do not include error bars in our experiment.Guidelines:• The answer NA means that the paper does not include experiments. 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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [No] Justification: We believe this work makes preliminary efforts towards understanding a potentially principled way of designing Transformer-like 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 no societal impact or why the paper does not address societal impact.Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

- 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.
- 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.
- 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.
questions: questions: 1. I cannot see why CRATE-N makes sense. If it is gradient ascent, how does it minimize $R^c$?
2. 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?
3. 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. 
4. Could you also use models that are not CRATE in the experiment of Section 5, and see if SRR can still predict their generalization?

**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: Limitations are not discussed. I encourage the authors to add a limitation section in the revision.

**Post-rebuttal note:** After discussing with my fellow reviewers and the AC, I raised my rating to 7.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.

[1] Yu et al. White-Box Transformers via Sparse Rate Reduction (2023)
questions: 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: limitations: The authors have partially discussed the limitations in the checklist. I encourage the authors to create a separate section in Appendix.
paper_summary: paper_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.
strength_weakness: strength_weakness: Despite the interesting observations, there seem to be disconnections here and there. 

## Section 4.1:



1. Lines 125 and 135-136 make conclusions based on how (6) is related to (7), but these conclusions are not justified. 

    _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._

* 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. 

    _When taking the gradient of R c to construct the MSSA operator, omitting its first-order term will produce a counterproductive effect._

* Maybe omitting the first-order term is not the (only) culprit, adding softmax could also play a role.
2. 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?

    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.

3. 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
    1. the original gradient update
    2. the update from 2nd order Taylor expansion, i.e., equation (6)
    3. the update from removing the second order term from (6)
    4. the update from further adding softmax, i.e., (7)

## Section 4.2:



1. 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.

    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?


## Section 5.6



1. 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.
2. 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.
questions: questions: See above
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
2. 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.
questions: questions: NA
limitations: limitations: NA

Comments: title: Thank you for the rebuttal
Thank you for the rebuttal
Response
comment: I thank the author for the interesting rebuttal. The results of using SSR to measure other transformers are interesting. I encourage the author to add these experiments to the camera-ready and explore more.
Thank you for acknowledging our responses and raising the score, and we are happy to see that your concerns are mostly addressed. To clarify more, we want to explore architectures that may outperform CRATE-C, the conceptual framework, and hopefully be on par with CRATE without bringing in new parameters that are hard to analyze. Efforts are made both theoretically (CRATE-N) and empirically (CRATE-T). Yes, we also find that CRATE-T and other variants (Table 1 in the supplemented pdf) may not outperform CRATE, and this further motivates us to perform deeper analysis on the relationship between SRR measure and performance, i.e., Section 5, which reveals the positive effects of SRR. Therefore, we believe that SRR could be useful for enhancing performance, but its utilization or its guidance to the model architecture needs to be further investigated. In particular, it is possible to find a better approximation or implementation of SRR in building transformer-like models, that could achieve a better trade-off between interpretability and performance than the current CRATE. This will be left for our future exploration.
Thank you for the comprehensive rebuttal, and I have read other reviews and responses as well. My original concerns were mostly addressed. I have adjusted my scores, but not very confidently since it seems a remaining limitation of the work is that CRATE-T, although motivated as a method that could outperform CRATE-C or CRATE by sacrificing interpretability (as in the response to W1 of reviewer RBPW), does not seem to outperform the original CRATE in the presented experimental results (e.g. in Table 2). This seems like a limitation as it means theoretical motivations were not precisely demonstrated in the experiments.
We're glad that the additional experiments address your concerns in Section 4.1 and we will make corresponding improvements in the revision.  As for concerns in Section 4.2,  although CRATE outperforms CRATE-C, we believe there is no need to introduce new parameters to do so. Optimizing the SRR objective in the forward pass could matter, but reducing the SRR measure as a whole in (10) could be more important to correlate with improved generalization, as demonstrated in Section 5. CRATE-N/T are just instantiations of this possibility, though motivated differently. Indeed, we need to investigate further the role of ISTA block in (3) to have a deeper understanding of the framework.
Thanks to the authors for their rebuttal!

**Section 4.1/1**

I agree that using softmax does not change the conclusion that update (7) is increasing the $R^c$ term (which is counterproductive) as shown in Figure (1). 

What I am not sure about is the apriori attribution of this counterproductive phenomenon to solely omitting first-order terms, as the paper writes on lines 135-136. Couldn't it be due to the softmax approximation? Do I miss something here?

**Section 4.1/2**

I see! So the thing that I was missing was the fact that $\lambda_i^k \geq 1 $ (which you mentioned on line 122), so one is always at the right half of the green curve. This could be highlighted better: e.g., you could make the left half of the green curve dashed, and/or reiterate the fact that $\lambda_i^k \geq 1 $ near line 140. 

**Section 4.1/3**

Thanks for providing the experiments. They make the picture clear, and therefore I would suggest adding/referring to them in section 4. 

**Section 4.2**

I think reviewers RBPW (W1), uGk9 (W1) and I share the same concern. It seems that CRATE-N and CRATE-T do not connect tightly with the findings in section 4.1. 

Perhaps the interesting finding is that of Table 2: Implementation-wise, one has to deviate from CRATE-C and use CRATE for maximized performance, and CRATE-N/-T serve as some sort of performance interpolations between CRATE-C and CRATE. It remains an open problem what the changes (from CRATE-C) do to the representation/optimization dynamics in terms of first principles.
We are glad that our response helps address your questions. Thank you for recognizing the additional experiments, and we will include these and more in the revision.
I thank the authors for the rebuttal. I like the additional experimental results presented in Figures 1 and 5 of the new pdf, and I encourage the authors to do more experiments even after the rebuttal period and include the new results in the new version. Most of my questions have been addressed, and I am inclined to accept this paper. I will discuss with my fellow reviewers and the AC and notify the authors if I raise my rating.