White-Box Transformers via Sparse Rate Reduction

In this paper, we contend that the objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a mixture of low-dimensional Gaussian distributions supported on incoherent subspaces. The quality of the final representation can be measured by a unified objective function called sparse rate reduction. From this perspective, popular deep networks such as transformers can be naturally viewed as realizing iterative schemes to optimize this objective incrementally. Particularly, we show that the standard transformer block can be derived from alternating optimization on complementary parts of this objective: the multi-head self-attention operator can be viewed as a gradient descent step to compress the token sets by minimizing their lossy coding rate, and the subsequent multi-layer perceptron can be viewed as attempting to sparsify the representation of the tokens. This leads to a family of white-box transformer-like deep network architectures which are mathematically fully interpretable. Despite their simplicity, experiments show that these networks indeed learn to optimize the designed objective: they compress and sparsify representations of large-scale real-world vision datasets such as ImageNet, and achieve performance very close to thoroughly engineered transformers such as ViT. Code is at \url{https://github.com/Ma-Lab-Berkeley/CRATE}.

Paper

Similar papers

Peer review

Reviewer tJjQ8/10 · confidence 4/52023-06-22

Summary

This paper proposes a optimization target called sparse rate reduction, which is built on previous optimization target called rate reduction [49]. By unrolling the iterative optimization process of sparse rate reduction into neural layers, a transformer-like architecture can be obtained. The derived white-box transformer-like architecture achieves similar performance with ViT.

Strengths

Overall the manuscript is well written and related works are properly cited and discussed, very insightful work. The idea is novel. This manuscript provides a significant extension to previous redunet [49]. The rate reduction is extended into sparse rate reduction, based on which transformer-like neural architectures can be derived by unrolling the iterative optimization process. The manuscript provides new insights concerning several important aspects of modern neural networks, i.e., score function is shown to be connected to self-attention and rate reduction under idealized token distribution. The results are promising, showing that the idea of white-box unrolling-based neural network design might be a possible alternative to current black-box design.

Weaknesses

The results on imagenet are very promising, while it would be more convincing if the proposed white-box architecture can achieve SOTA performance under fair comparison. Although overall the manuscript is well written, some sentences are too concise and a little confusing. I suggest the authors can go over the whole manuscript and improve the text for general readers. 1) L143 and footnote 4, I think the explaination here is too concise, the footnote confuses me. What is the separation and the mathematical roles? I cannot find related contents elsewhere. 2) L185 and footnote 6, can you explain more about the content in footnote 6? What is the rigorous math here? 3) I think the motivation of sparse coding (the L0) term in eq. 1 can be further clarified. I can understand the motivation of rate reduction [49], but I think it is not clear why we introduce sparse coding here. Further, is L0 norm the best implementation of sparse coding here? I hope the authors can discuss more about the designing principle of those optimization metrics or targets like eq. 1. Following previous question, I think the optimization metric is closely related to specific task. In my opinion, rate reduction has a natural connection to classification problem. Can the authors comment on the connection between optimization metric and specific tasks? If we are considering object detection or more complex problem like image reconstruction, what is the general principle that can guide the design of optimization metrics?

Questions

see above.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

see above.

Reviewer YxEu6/10 · confidence 3/52023-07-01

Summary

The authors propose an interpretation of the transformer architecture wherein the component blocks may be interpreted as unrolled optimization steps: * Multi-Head Self-Attention (MHSA) is said to be approximately the same as Multi-Head Subspace Self-Attention, which is an unrolled optimization of the following objective: * $ \sum_{k=1}^K R ( \mathbf{U}_k^* \mathbf{Z})$ (see eq. 8), where $R$ is an estimator for the "lossy coding rate" $R(\mathbf{Z}) = \frac{1}{2} \text{logdet} ( \mathbf{I} + \frac{d}{N \epsilon^2} \mathbf{Z} \mathbf{Z}^* )$ * Multi-Layer Perceptrons are said to be approximately the same as Iterative Shrinkage-Thresholding Algorithms (ISTA) which are an unrolled optimization of the following objective: * $\lambda ||\mathbf{Z}||_1 + ||\mathbf{Z}^{l+1/2} - \mathbf{D} \mathbf{Z}||^2_F$ where $\mathbf{Z}$ are activations, $\mathbf{Z}^{l+1/2} = \mathbf{D} \mathbf{Z}^{l+1}$ and this is justified as being a relaxed LASSO objective that will sparsify the representation of $\mathbf{Z}$ The relationship between the architecture and the objective is quite involved and I cannot summarise it here. As a consequence of constructing the network in this way, the authors are able to track both of these objectives as a way to gain insight how a network is operating; this is the "white-box" of the title. Specifically, they track the sparse coding $R$ mentioned above and the sparsity at each layer in a network. During training it is then possible to observe the sparse coding rate and sparsity both decrease with the layer index during training. In addition, the network modifications required to match the theory do not appear to reduce performance on ImageNet compared to a similarly sized Vision Transformer.

Strengths

Understanding the performance and popularity of self-attention in deep learning is a valuable goal and finding a theoretical motivation for this specific architecture could be very useful to anyone using or developing transformers. If this goal is achieved by this paper then it is a significant work. The strengths of this paper are therefore mainly found in Figures 3 and 4, which show the statistics $R^c$ and sparsity. Figure 3 compares train versus validation while Figure 4 shows how these statistics change during training. It is interesting to observe these values decreasing as the signal travels through the network as this approximately matches what is predicted by the theory. These are the original empirical observations of the paper and they support the paper's claims. The motivation and experimental results are both clearly stated. Section 1 provides a reasoned argument for why learning representations that are interpretable would be valuable in current deep learning practice, and that it is lacking in current architectures, either diffusion or autoregressive models. The results in Table 1 also address concerns that this method would reduce performance, it looks like performance is mostly maintained, which is promising.

Weaknesses

Addressing following weaknesses would substantially improve the paper * Legibility of the derivation: the derivations in Section 2 are dense and difficult to follow * The sparse rate coding function $R^c$ is key but is only introduced in Section 2.3, introducing it earlier and giving the reader an idea of how it relates to the distribution of $\mathbf{Z}$ would make it much easier to understand. For example, noting it's relationship to the entropy of a multivariate normal, how it will grow with the variance of $\mathbf{Z}$ etc * At the end of Section 2.2 the reader has just finished reading a derivation of "Self-Attention via Denoising Tokens Toward Multiple Subspaces" and then immediately afterwards is faced with "Self-Attention via Compressing Token Sets through Optimizing Rate Reduction". I am still confused as to which derivation I'm supposed to pay attention to for understanding the MSSA block. If both derive the same MSSA block then include only one and put the other in the Appendix, this will also help free up space for more experimental results * Do not rely on terminology from previous papers that are not common knowledge, for example if you agree that most readers will not understand the usage of "incoherent" without explanation, do not use it like that. Same for "linearized and compact". * On line 123 there's a reference to (8), which appears to be the ImageNet paper from 2008 and I could not find a formal definition of the lossy coding rate in it * The theory demonstrates how each layer could be performing an unrolled optimization step but it does not explain why this is beneficial to the overall problem of learning a function, such as predicting the class label using the entire network; in fact the target class labels are not present in the notation. I believe the entire network was optimized according to the cross-entropy loss, so where is that in Section 2? * Additional experimental results would be worthwhile: * Table 1 does not contain results mentioned on lines 354-355 showing the scaling performance of ViT. Even if these results are from other work, including it with citation could be valuable in context * It is known that ViT architectures underperform on ImageNet versus larger datasets. Unfortunately, it would be valuable to see results training on Imagenet-21k, as in the original ViT paper, but this may be beyond the authors resource capacity * Computing $R^c$ and sparsity with depth in public pretrained models would also be useful, either to demonstrate that these models fail to minimize these implied objectives or to demonstrate that this is why transformers work * Comparing these statistics to activation norms would be a good comparison

Questions

What experiment could you plan that would disprove the theory in Section 2? What architectural modification would not allow a low sparse coding rate or sparsity but would still allow similar performance when trained? What do these results imply for transformer design, can you predict anything that is obviously incorrect or valuable just from these results? Can you design a new layer that performs an additional form of unrolled optimization that is also useful? Is there a regime of transformer operation that completely fails where this theory would provide a useful insight? I think this direction of research is valuable and if I could understand precisely why the white-box observations in Figure 3 and 4 are extremely valuable I would increase my rating.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitation of trying to explain why a given deep learning architecture works is: 1. A method becomes popular because it is surprisingly effecitive 2. A theory is constructed to explain why it works 3. The theory demands some small change to the network which only slightly reduces performance 4. The theory does not immediately permit an improvement to the method that is valuable I think this is why Figures 3 and 4 are key. If the statistics shown in Figures 3 and 4 were extremely valuable to understanding transformer training dynamics then it would be obvious that this work is significant. It could be that they are, but I do not see this argument made clearly enough in the paper. Alternatively, maybe there is some architectural or training improvement that this theory implies that would be extremely valuable.

Reviewer vvYz7/10 · confidence 3/52023-07-05

Summary

This paper proposes to structure a classification pipeline based on Transformer networks using precisely defined mathematical operators that are designed to perform a gradient step to minimize a well defined objective, e.g. Lasso objective sparse representation, or maximizing auto-correlation between "noisy" and "denoised" tokens. As such, it not only propose a new architecture with some shared weights, but it also proposes an interpretation for the role of each block to achieve the goal of representation learning. The pipeline is tested on popular Image Classification benchmark with end-to-end learning (ImageNet) as well as transfer learning (pre-trained on ImageNet, finetuned on CIFAR-10/100, Oxford flowers and pets.

Strengths

The paper really proposes a new architecture based on an intuition " the objective of representation learning is to compress and transform the distribution of the data towards a mixture of low-dimensional Gaussian distributions supported on incoherent subspaces ". They propose an architecture which resembles the visual transformers with a strong effort of modeling, i.e. trying to affect an objective to the usual transformers blocks (i.e. multi-head attention). The proposed architecture is conceptually significantly simpler than ViT for example, and the performance drop is arguably very slight (-1% top1 on ImageNet). A posteriori analysis before and after training (e.g. Figure 4) seem to validate the intuition of the authors for the coding rate aspect (it seems not so clear for sparsity). This modelisation work in a tough work, and this paper is a significant contribution.

Weaknesses

To me there is a caveat about the sparse coding hypothesis. Line 252: " In our implementation, motivated by Sun et al. [29] and Zarka et al. [31], we also add a non-negative constraint to Z^{l+1} " This non-negative constraint is not just a detail, compared to a ReLU, the soft-thresholding reduced and possibly zeros the coding coefficient but it preserves the sign, i.e. the phase. The ReLU collapses the sign, i.e. the phase. This "detail" is also under-estimated in Sun et al. [29] and Zarka et al. [31] . I suggest the authors to read the follow up work by Guth et al https://arxiv.org/pdf/2110.05283.pdf especially section 4 Phase Collapse Versus Amplitude Reduction. This might give further intuitions in this foggy world of modeling neural networks. I'm not claiming that phase collapse is the good interpretation for this block, but for sure SoftShrink vs ReLU, i.e. sparse code vs non-negative sparse code is probably not a detail.

Questions

How much does the performance degrade when replacing ReLU by softshrink ? Do you have ideas / remarks / intuitions on the possible importance of this non-negativity constraint ? How does it articulate with the mixture of Gaussian / low-dimensioal subspace hypothesis ? Gaussians do no care about the sign, isn't it ? Note that with non-negativity constraints , a subspace becomes a cone. Why would it be important to have a cone rather that a subspace ?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Apart the aspect mentioned below, this is a really strong work.

Reviewer 13ev5/10 · confidence 2/52023-07-10

Summary

The authors propose a novel theoretical framework which shows that the popular transformer can be motivated by Maximizing Rate Reduction. The key idea of this work follows previous information-gain framework, mostly ReduNet, but with more special careful treatment on connections to transformer architecture design. By a few approximations, the authors show that maximizing rate reduction with sparsity constraint indeed derives a transformer-like deep structure. The derived white-box transformer-like architectures are verified on multiple datasets.

Strengths

1) The paper is overall well written and easy to follow. Related Work section includes comprehensive surveys. Formulas are clearly explained. Experiments come with detailed settings. 2) The idea is novel and interesting. Especially, deriving multi-head attention from Maximizing Code Reduction is novel. 3) The proposed white-box architecture is verified on real-world datasets such as ImageNet-1k and compared to ViT models.

Weaknesses

1) Although the paper claims that the proposed white-box model is competitive to ViT models, the numerical results seem not strong. On ImageNet-1k, the proposed model is clearly underperforming even with larger number of parameters. 2) The real power of ViTs is on high accuracy regime, where the model size is large. The authors only consider small model regime with low to medium accuracy, which lacks of convincing.

Questions

1) Please consider larger model with ImageNet-1k top-1 accuracy above 80.0%

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Reviewer AAKz10/10 · confidence 4/52023-07-12

Summary

This paper provides an interesting claim "the standard transformer block can be derived from alternating optimization on complementary parts of this objective: the multi-head self-attention operator can be viewed as a gradient descent step to compress the token sets by minimizing their lossy coding rate, and the subsequent multi-layer perceptron can be viewed as attempting to sparsify the representation of the tokens". This results in white-box transformer-like deep network architectures which are mathematically fully interpretable.

Strengths

- Very interesting claim about transformer. - Extensive experiments to verify the effectiveness of the proposed method.

Weaknesses

I tried and failed to find any weaknesses. I really like such work on interpretable neural networks.

Questions

No.

Rating

10: Award quality: Technically flawless paper with groundbreaking impact, with exceptionally strong evaluation, reproducibility, and resources, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

No.

Reviewer YxEu2023-08-11

No argument justifying award quality score of 10

I'm sorry to point this out but in order to justify a score of 10, it is necessary to write more than four sentences. The assertion that there are no weaknesses to the work does not seem reasonable, there are certainly weaknesses as pointed out by other reviewers. I would be very interested to read a convincing argument as to why this paper is award quality and I hope you will be able to revise this review to provide it. If not I hope an unjustified appraisal of award quality does not affect the final decision on this paper's acceptance.

Reviewer YxEu2023-08-11

Falsification Results

In reply to this comment: > We find that without our white-box design, the vanilla ViT does not optimize our proposed sparse rate reduction objective. This contrasts with the results shown in Figures 3 and 4 of the work, wherein we see that the $R^c$ and sparsity value decrease layerwise for CRATE, in accordance with our theory. Since CRATE is an architectural modification of the vanilla ViT, we think this presents a compelling answer to your questions about an architectural modification that achieves similar performance without the same white-box characteristics, as well as a compelling experiment that could falsify our theory in Section 2 (and fails to do so). I don't understand the results presented in Figure 1 of your pdf rebuttal: 1. The sparsity of a vanilla ViT should not be 1.0, in that case no activations would ever be sufficiently negative entering the GeLU activation functions, which would mean the model is failing to learn nonlinear functions. Specifically, [other papers][lazy] have demonstrated activation sparsity of 6.3% in vision transformers recently. 2. While the coding rate decreases for CRATE as the layer index increases, I don't see this as being a very conclusive result, it only decreases 27% through the entire network, while ViT-B decreases 15% from layers 4 to 8. I apologise for the confusion because I did not state what I thought your hypothesis was so that makes it difficult to talk about falsification. The hypothesis I had in mind was, "sparse rate reduction is sufficient for learning useful functions in deep networks". In that case, I don't see this as falsifying anything, if sparse rate reduction is critical for your network to learn then what you need to is ablate those capabilities from the model and show that it is no longer able to learn. In fact, if I accept the statement above that the ViT has no capability for sparse rate reduction and is still able to outperform CRATE then I can only conclude the sparse rate reduction in CRATE is irrelevant. I would attribute the performance to the architectural similarity to a transformer and the effectiveness of contemporary minibatch SGD on a cross-entropy objective. The derivation is interesting and the results matching contemporary networks are promising but I don't see what the "white box" buys you. Specifically: 1. It's not significantly correlated with performance, your results comparing to ViT demonstrate this 2. It doesn't provide a significant benefit in architectural design beyond prescribing a block that is similar to a transformer 3. It doesn't change how models are trained in any significant way (unless the paper fails to mention that this model converges in a significantly different way to a transformer) 4. It encourages sparsity but the computation benefit isn't explored in the paper, nor is it demonstrated that this sparsity is significantly lower than it is now known to be in transformers. For example, ["The Lazy Neuron Phenomenon"][lazy] demonstrates 6.3% nonzero entries in ViT-B16 5. Interpretability is a fuzzy concept, but I don't see any experiments in the paper aimed at interpreting what the network is doing based on the sparse rate reduction metrics, the experiments simply observe the metrics decrease while the network learns [lazy]: https://arxiv.org/abs/2210.06313

Authorsrebuttal2023-08-12

Discussion with Reviewer YxEu (Part 1)

We are grateful for you engaging with our rebuttal further, and for your critical perspective on the work, which will no doubt improve it. Thank you also for pointing out that **‘[t]he derivation is interesting and the results matching contemporary networks are promising.’** ### Interpreting the requested experiments on Figure 3/4 on public ViTs Our results are an accurate reflection of the experiment that you suggested (e.g., computing $R^c$ and sparsity as a function of depth). The last two columns of the first two rows of Figure 1 in the .pdf rebuttal evaluate the sparsity of the tokens after the second block of each transformer layer, $z\_{\\ell} = \\mathrm{MLP}(\\mathrm{LN}(z^{\\prime}\_{\ell})) + z^{\\prime}\_{\\ell}$, as defined in Eq. (3) of the ViT paper [1]), which make this a consistent comparison with how we evaluate the sparsity of the second block (i.e., ISTA block) of each CRATE layer. We applied the original weights and architecture of the public pre-trained ViT model from the `timm` package. Compared with the other paper [2] you mentioned, [2] evaluated the sparsity of the hidden layer output of the MLP, which is different from ours. Meanwhile, [2] replaced the GeLU activation with ReLU activation in the MLP layer. We did not find public checkpoints for the models in [2]. Similarly for the $R^c$ results, although we can agree that a subjective interpretation is possible, we think it is unambiguous that in a network derived from unrolled optimization, it “optimizes” the objective if the objective trends in the appropriate direction on average. The result shows that for the compression part of the objective, this is true of CRATE, and not true of ViT; the CRATE-S model’s $R^c$ term is reduced by about 30% over the course of forward propagation, whereas the analogous terms for ViT-S and ViT-B increase by about 50%. If you would like to see additional comparisons here, we are happy to run them and report the results. **We want to state clearly that the results we are reporting are accurate**; we are only providing our interpretations of the results of the experiments you suggested. But let us emphasize that evaluating the $R^c$ and sparsity metrics in the parts of the network we have is most reasonable for CRATE precisely because _we have designed the network to learn a representation that has these characteristics_. Your suggested experiments demonstrate that the token embeddings of the ViT – analogous to what we evaluated in Figures 3 and 4 – do not have these same properties. However, this does _not_ imply that the ViT does not learn low-dimensional or parsimonious (e.g., compressed and sparse) representations of the data. Rather, it implies that the ViT’s learned representations are less accessible, and thus harder to evaluate, due to its parameter-redundant black-box design. This is a key benefit of our derivation, and the simplified white-box architecture of CRATE: the places where the representations are transformed to standard forms (axis-aligned, hence sparse, orthogonal subspaces) are completely exposed to the network architect, removing any ambiguity in measuring these quantities. We believe these insights present an excellent opportunity for follow-up work to better understand the ViT, as well, but this is firmly out of scope of the present submission. [1]: A. Dosovitskiy et al., “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale,” in International Conference on Learning Representations, 2021. [2]: Z. Li et al., “The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers,” in The Eleventh International Conference on Learning Representations, 2023.

Authorsrebuttal2023-08-12

Discussion with Reviewer YxEu (Part 2)

### Our work’s hypotheses Thank you for clearly stating your thinking on this point. We think you may be misunderstanding our work’s principal experimental hypothesis: "sparse rate reduction is sufficient for learning useful functions in deep networks" seems to us to be a misinterpretation. Let us also mention in this connection that your assertion > if I accept the statement above that the ViT has no capability for sparse rate reduction and is still able to outperform CRATE then I can only conclude the sparse rate reduction in CRATE is irrelevant. does not seem to be logically sound: **we have not claimed anywhere that using the sparse rate reduction is necessary to construct high performance deep models**. In fact, its role in our derivation is quite the opposite: its use in the design of the architecture promotes the learning of mathematically-interpretable representations of the data in the network. Please see the discussion of “what white-box buys us” for more on this point. To clearly state our central hypothesis, let us reiterate our primary motivations, which were written in Section 2.1. Our goal is to design a network architecture that transforms the data to a mixture of nearly-orthogonal axis-aligned subspaces, the optimizers of the sparse rate reduction [3]. We thus obtain our architecture from unrolled optimization on the sparse rate reduction, then learn its parameters with backpropagation, since the structure of the data distribution is unknown. In particular, since we are learning these parameters, there is no guarantee that the resulting network will optimize a sparse rate reduction objective for the data distribution. This leads precisely to **our main hypothesis, that *it is possible to train a transformer-like architecture (i.e., CRATE) to simultaneously achieve high accuracy at scale and optimize the sparse rate reduction***. The results in the left panel of Figure 4 demonstrate that at random initialization, the CRATE-Small model does not optimize the sparse rate reduction for the data distribution – only through learning does the network optimize the sparse rate reduction for the data distribution. Note that in this line of reasoning, the *goal* is to obtain a useful representation of the data. We have argued in the introduction of the submission why this goal is valuable and a central ‘grail’ for learning. **The goal is not only to obtain a network with high performance; the goal is to obtain a white-box network which learns useful representations.** [3]: Y. Yu et al., "Learning diverse and discriminative representations via the principle of maximal coding rate reduction." Advances in Neural Information Processing Systems 33 (2020).

Authorsrebuttal2023-08-12

Discussion with Reviewer YxEu (Part 3)

### What white-box buys us We would like to push back on your characterization of our white-box model. In our work, a “white-box model” can be thought of as a model whose architecture and parameters are derived mathematically from first principles, in a manner where the data distribution plays a central role. In this view, your first assertion is self-evident: the most natural white-box model for representation learning would be sparse coding of the data, possibly in a learnable signal dictionary, which gives a mathematically-interpretable and practically-robust model with performance that unfortunately cannot match that of modern deep learning architectures. Our contribution is to present a white-box derivation of a transformer-like architecture that is simultaneously highly performant. We truly believe there is significant novelty in this contribution: despite notable efforts from the theoretical community to suggest possible interpretations for the self-attention operation in transformers (e.g., summarized in [4]), a holistic and practically-verified interpretation for an entire transformer-like block (i.e., both the self-attention operation and MLP) has not been proposed before our work. In response to your second point, we would like to reiterate that as we wrote in the rebuttal, there are in fact **concrete practical implications** of our work for standard transformers: specifically, that **the QKV matrices in self-attention layers of ViT are redundant, and can be combined to save almost a factor of 4 in the overall parameter count**, with only a minor performance hit that can surely be reduced further with additional engineering work. Regarding your remaining points, we already mentioned the interpretability experiments we conducted in the submission in our rebuttal “**General Comments**” (e.g., visualizing the learned dictionaries and subspaces of CRATE in Figures 6 through 11 in the appendix); we think studying the computational benefit of sparsity based on our results is an interesting direction for future work, but firmly out-of-scope for the present work. We appreciate from your response here and to `AAKz` that you harbor some skepticism of research on the “model-centric” understanding of deep networks – your valuation of our work seems to be primarily a function of the extent to which such work directly implies improvements to specific metrics in practice. Consider that, if some of the significant methodological innovations in deep learning from the last five years were subjected to the same standard, they would have been dismissed – for example, diffusion models were not demonstrated to have sample quality anywhere close to the state-of-the-art GANs of the time [5]. The important aspect of these works was their conceptual insight that pointed the community towards more principled approaches and led to tremendous performance gains in the long run. We believe that CRATE has similar potential for future developments – it becomes possible to realize other novel improvements not just through empirical design, but also by using the guidance of principles from optimization and compression through the white-box approach. [4]: R. Vidal, “Attention: Self-Expression Is All You Need,” Sep. 29, 2021. Accessed: Apr. 05, 2022. [5]: J. Sohl-Dickstein et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

Reviewer YxEu2023-08-13

Thank you for this direct and comprehensive response, I apologise for the bluntness of my previous comments, I was aiming to be brief. > We appreciate from your response here and to AAKz that you harbor some skepticism of research on the “model-centric” understanding of deep networks – your valuation of our work seems to be primarily a function of the extent to which such work directly implies improvements to specific metrics in practice. Consider that, if some of the significant methodological innovations in deep learning from the last five years were subjected to the same standard, they would have been dismissed – for example, diffusion models were not demonstrated to have sample quality anywhere close to the state-of-the-art GANs of the time [5]. I agree that this is not a standard we should hold research to, and I understand that asking for state of the art results is not reasonable in this case. My goal was trying to make the correspondence between the theoretical results and experiments clear, I didn't have a good frame of reference of whether a $R^c$ value of 1600 means anything compared to 900, or whether sparsity of 0.2 is a good result for your model specifically. I can see that those numbers go down during training, which is what we should hope to see from the derivation of the unrolled optimization but it did not seem conclusive to me. # Figure 3/4 Sparsity Result To be precise, my misunderstanding is the sparsity you're measuring is at the output of a residual MLP and in [2] the authors only measure sparsity at the output of activation functions (ie no residual connection is present). Looking at the ISTA block in Figure 2, the output is an activation function. In my opinion a more useful metric in Figure 1 of the pdf rebuttal would be the output of the activation function of the MLP, ie: $$ z_l = MLP(LN(z_l')) + z_l' $$ $MLP(LN(z_l'))$ is functionally more similar to the output of the ISTA block and much more likely to be sparse, because it is the output of an activation function. It is no surprise for $z_l$ to exhibit no sparsity because $z_l'$ is the output of the self-attention block. Whether this is a relevant comparison to CRATE is another question, but I think that figure would have made much more sense if it had represented the sparsity present in the ViT. # Necessary vs Sufficient My statement was "sparse rate reduction is sufficient for learning useful functions in deep networks" and you reject "using the sparse rate reduction is necessary to construct high performance deep models". These statements are not saying the same thing. In fact I think we agree on what the hypothesis of your work is, ie that CRATE is sufficient to learn useful deep representations. # What white-box buys us I agree that the following contributions are valuable: 1. The concrete example given: "the QKV matrices in self-attention layers of ViT are redundant, and can be combined to save almost a factor of 4 in the overall parameter count" 2. Visualizations in Figures 6 to 11 in the Appendices could be informative 3. Building a network from "axis-aligned, hence sparse, orthogonal subspaces" is likely to open up other valuable lines of inquiry Your response here, and further reviewing the Appendices, have made that clear. I will update my review based on this. [2]: Z. Li et al., “The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers,” in The Eleventh International Conference on Learning Representations, 2023. [5]: J. Sohl-Dickstein et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

Authorsrebuttal2023-08-14

Discussion with Reviewer YxEu

Thank you for engaging with our rebuttal further and raising your score after reviewing our response. We are grateful for your valuable feedback on our work. > To be precise, my misunderstanding is the sparsity you're measuring is at the output of a residual MLP and in [2] the authors only measure sparsity at the output of activation functions (i.e., no residual connection is present). - For ViT, we measured the sparsity at the output of the residual MLP block, i.e., $z\_{\\ell} = \\mathrm{MLP}(\\mathrm{LN}(z^{\\prime}\_{\\ell})) + z^{\\prime}\_{\\ell}$. - For our proposed CRATE model, we measured the sparsity at the output of the ISTA block. We would like to clarify that, in CRATE, the residual connection appears in the ISTA block before the activation: see Eq.(16) in the main body for the precise structure. > In my opinion a more useful metric in Figure 1 of the pdf rebuttal would be the output of the activation function of the MLP. We agree that measuring the sparsity of $\\mathrm{MLP}(\\mathrm{LN}(z^{\\prime}\_{\\ell}))$ is another useful metric in addition to our existing measurements. Note that in the ViT, we have $\\mathrm{MLP}(\\mathrm{LN}(z^{\\prime}\_{\\ell})) = W\_2 \\sigma(W\_1 \\mathrm{LN}(z\_\\ell^\\prime))$, where $\sigma$ is the activation function in the MLP block (in the `timm` ViT this is GeLU while in [2] this is ReLU), so the MLP output is not actually the output of an activation function. We have visualized the sparsity of $\\mathrm{MLP}(\\mathrm{LN}(z^{\\prime}\_{\\ell}))$ for ViT and did not observe clear sparsity patterns, e.g., the ratio of non-sparse elements is larger than 90% for ViT-Base. We also measured the sparsity of $\\sigma(W\_1 \\mathrm{LN}(z^{\\prime}\_{\\ell}))$, the hidden layer activations; they are indeed sparse in a manner consistent with [2]. We hope this clarification makes clear our point from the last paragraph of Part 1 of our previous response about the fact that token-level representations in CRATE are sparse, whereas for ViT, what is sparse is a less-exposed hidden layer output in the MLP. We thank you for your suggestion here; we will add these results and this discussion to our camera-ready version. We thank you again for your thoughtful review and valuable feedback! Please let us know if you have any other questions or comments during the discussion period. [2]: Z. Li et al., “The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers,” in The Eleventh International Conference on Learning Representations, 2023.

Reviewer vvYz2023-08-13

Response to rebuttal

Thank you for this interesting discussion, happy to read that you have already lines of explanation for the non-negativity constraint, look forward reading the camera ready version.

Authorsrebuttal2023-08-14

Response to Reviewer vvYz

Thank you again for thoroughly reviewing our manuscript and response. We are grateful for your valuable feedback on our work, which will no doubt improve it. Please let us know if you have any other questions or comments during the discussion period.

Authorsrebuttal2023-08-14

Response to Reviewer tJjQ

Thank you again for thoroughly reviewing our manuscript and response and raising your score. We are grateful for your valuable feedback on our work, which will no doubt improve it. Please let us know if you have any other questions or comments during the discussion period.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC