White-Box Transformers via Sparse Rate Reduction
Authors: Yaodong Yu
Sam Buchanan
Druv Pai
Tianzhe Chu
Ziyang Wu
Shengbang Tong
Benjamin David Haeffele
Yi Ma
Venue: NeurIPS 2023 poster (2023)
Keywords: white-box deep neural networks
representation learning
transformer
sparse coding
TL;DR: We develop white-box transformer-like deep network architectures which are mathematically interpretable and achieve performance very close to ViTs.

Abstract: 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 https://github.com/Ma-Lab-Berkeley/CRATE.

White-Box Transformers via Sparse Rate Reduction
Yaodong Yu' Sam Buchanan Druv Pai
Tianzhe Chu'
Ziyang Wu'
Shengbang Tong'
Benjamin D. Haeffele
Yi Ma-
University of California, Berkeley
2TTIC Johns Hopkins University
Abstract
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 unified objective function called spurse rale reduction. From this perspective, popular deep networks such as transformers can be naturally viewed as realizing iterative schemes optimize this objective incrementally. Particularly, we show that the standard transformer block can derived from alternating optimization on complementary parts this objective: the multi-head self-attention operator can be viewed 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-lie deep network architectures which 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 VIT. Code is at https: htaps:///ha thub. Com//ta-Lab-Berkeeeeee Berkeley/CRATE
Introduction
In recent years, deep learning has seen tremendous empirical success in processing massive amounts of high-dimensional and multi-modal data. Much of this success is owed effective learning of the data distribution and then transforming the distribution to a parsimonious, i.e. structured and compact, representation [39, 50, 52, 62], which facilitates many downstream tasks (e.g., in vision, classification [23, 40], recognition and segmentation [25, 38, 77], and generation [31, 65, 66]). To this end, many models and methods have been proposed and practiced, each with its own strengths and limitations. Here, we give several popular methods a brief accounting as context for a complete understanding and unification that we seek in this work.
Transformer models and self-attention. Transformers [28] are one of the latest popular models for learning representation for high- dinensinnns structured data, such text [28, 30, 37], images [40, 75], and other types of signals [48, 57]. After the first block, which converts each data point (such text corpus image) into on the token sets, in a medium-agnostic manner [28, 40]. cornerstone of the transformer model is the so-called self-attention layer, which exploits the statistical correlations among the sequence of tokens to refine the token representation. Transformers have been highly successful in learning compact representations that perform well on many downstream tasks. Yet the transformer network
sequence of tokens, further processing is performed
yyyu, (ema)aacaeee.eee..eeooo www.eley..eeoooo..ooooo......eoe edu, (druvpai, chitpp,,,chuth,,,, zywu, tsb)/berkeley..de edu 2 sam@ttic, edu bhaeffele@jhu.eed edu
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Mut-HHead Supspace Sell-Attention (KBBA)
Spense Codingg Provinal Step
(187A)
compression
spersification
z
2
The 'main loop' of the CRATE white-box deep network design. After encoding input data X as sequence of tokens zº. CRATE constructs deep network that transforms the data to a canonical configuration of low-dimensional subspaces by successive compression against local model for the distribution, generating z441/2, and sparsification against a global dictionary, generating z++1, Repeatedly stacking these blocks and training the model parameters via backpropagation yields a powerful and interpretable representation of data.
Figure 1:
architecture is empirically designed and lacks rigorous mathematical interpretation. In fact, the output of the attention layer itself has several competing interpretations [68, 78]. As a result, the statistical and geometric relationship between the data distribution and the final representation learned by a transformer largely remains a mysterious black box.
Diffusion models and denoising, Diffusion models [22, 34, 41, 43, 44] have recently become a popular method for learning the data distribution, particularly for generative tasks and natural image data which are highly structured but notoriously difficult to effectively model [3. 5]. The core concept of diffusion models is to start with features sampled from a Gaussian noise distribution (or some other standard template) and iteratively denoise and deform the feature distribution until it converges to the original data distribution. This process is computationally intractable if modeled in just one step [61], so it is typically broken into multiple incremental steps. The key to each step is the so-called score function, or equivalently [13] an estimate for the "optimal denoising function""; in practice this function is modeled using a generic black-box deep network. Diffusion models have shown effectiveness at learning and sampling from the data distribution [56, 60, 65]. However, despite some recent efforts [81], they generally do not establish any clear correspondence between the initial features and data samples. Hence, diffusion models themselves do not offer a parsimonious or interpretable representation of the data distribution.
Structure-seeking models and rate reduction. In both of the previous two methods, the represen- tations were constructed implicitly as a byproduct of solving a downstream task (e.g.. classification or generation/sampling) using deep networks. However, one can also explicitly learn a representation of the data distribution as a task and of itself; this is most commonly done by trying to identify and represent low-dimensional structures in the input data. Classical examples of this paradigm include model-based approaches such as sparse coding [2, 29] and dictionary learning [17, 21, 47),, out of which grew early attempts at designing and interpreting deep network architectures [18, 32]. More
recent approaches build instead from a model-free perspective, where one learns a representation through a sufficiently-iformaaiis pretext task (such compressing similar and separating dissimilar data in contrastive learning [45, 69, 80], or maximizing the information gain in the class of maximal coding rate reduction methods [6, 46, 55]). Compared to black-box deep learning approaches, both model-basee and model-free representation learning schemes have the advantage of being more interpretable: they allow users to explicitly design desired properties of the learned representation [46, 55, 63). Furthermore, they allow users to construct new white- box nawaadd-constructed deep network architectures [11, 55, 59] by unrolling the optimization strategy for the representation learning objective, such that each layer of the constructed network implements an iteration of the optimization algorithm [11, 53, 55]. Several recent works [71, 74, 76] consider the connections between transformer architectures [28] and unrolled optimization. Unfortunately, in this paradigm, if the desired properties are narrowly defined, may be difficult to achieve good practical performance on large real-world datasets.
Our contributions, and outline of this work. In this work, we aim to remedy the limitations of these existing methods with more unified framework for designing transformer-lik network architectures that leads to both mathematical interpretability and good practical performance. To this end, we propose to learn a sequence of incremental mappings to obtain a most compressed and sparse representation for the input data (or their token sets) that optimizes a unified objective function known as the sparse rate reduction, specified later in (1). The goal of the mapping is illustrated in Figure Within this framework, we unify the above three seemingly disparate approaches and show that transformer- like deep network layers can be naturally derived from unrolling iterative
2

---

optimization schemes to incrementally optimize the sparse rate reduction objective. In particular, our contributions and outline of the paper are as follows:
In Section 2.2 we show, using an idealized model for the token distribution, that if one iteratively denoises the tokens towards a family of low-dimensional subspaces, the associated score function assumes an explicit form similar to a self-attention operator seen in transformers.
In Section 2.3 we derive the multi-head self-attention layer as unrolled gradient descent step to minimize the lossy coding rate part of the rate reduction, showing another interpretation of the self-attention layer as compressing the token representation.
In Section 2.4 we show that the multi-layer perceptron which immediately follows the multi- head self-attention in transformer blocks can be interpreted as (and replaced by) a layer which incrementally optimizes the remaining part of the sparse rate reduction objective by constructing a sparse coding of the token representations.
In Section 2.5 we use this understanding to create a new white-box (fully mathematically in- terpretable) transformer architecture called CRATE (i.e., Coding RAte reduction TransforEEr), where each layer performs a single step of an alternating minimization algorithm to optimize the sparse rate reduction objective.
Hence, within our framework, the learning objective function, the deep learning architecture, and the final learned representation all become white boxes that are fully mathematically interpretable. As the experiments in Section 3 show, the CRATE networks, despite being simple, can already learn the desired compressed and sparse representations on large-scale real-world datasets and achieve performance par with much more heavily engincered transformer networks (such VIT) wide variety of tasks (c.g.. classification and transfer learning).
2 Technical Approach and Justification
2.1 Objective and Approach
We consider general learning setup associated with real-world signals. We have some random variable X = - #N| E DxN which is our data source; each æ; E RR is interpreted as a token', and the æ;'s 's may have arbitrary correlation structures. We use Z -I E dR to denote the random variable which defines our representations. Each 2; E R is the representation of the corresponding token æ;, We are given B 2 i.i.d. samples XI,.... Xn ~ X, whose tokens are Fi.b- The representations of our samples are denoted Zi,... Zn ~ Z, and those of our tokens are Zi,b. Finally, for given network, we use z' to denote the output of the first f layers when given X as input. Correspondingly, the sample outputs are Z and the token outputs are z1.3-
Objective for learning structured and compact representation. Following the framework of rate reduction [55], we contend that the goal of representation learning is to find a feature mapping f: X E D + Z RR which transforms input data X E DxN with potentially nonlinear and molti-modal distribution to a (piecewise) linearized and compact feature representation Z R While the joint distribution of tokens (z.) in Z may be sophisticated (and task- specific), we further contend that it is reasonable and practical to require that the target marginal distribution of individual tokens z; should be highly compressed and structured, amenable for compact coding. Particularly, we require the distribution to be mixture of din-dimensional (say K) Gaussian distributions, such that the h Gaussian has mean 0 E Rd. covariance Σ₁ > 0 E d and support spanned by the orthonormal basis U2 E dx We denote UIK] ( to be the set of bases of all Gaussians. Hence to maximize the information gain [62] for the final token representation, we wish to maximize the rate reduction [6, 46] of the tokens, i.e., max z AR( Z; UKK) R R(()=
Ro Z: UiK). where R and R are estimates of lossy coding rates to be formally defined in (7) and (8). This also promotes token representations 2; from different Gaussians be incoherent [46].
Since rate reduction is an intrinsic measure of goodness for the representation, it is invariant to arbitrary rotations of the representations. Therefore, to ensure the final representations are amenable to more compact coding, we would like to transform the representations (and their supporting subspaces) so that they become sparse with respect to the standard coordinates of the resulting
For language transformers, tokens roughly correspond to words [28], while for vision transformers, tokens correspond to image patches [40].
3

---

representation space. The combined rate reduction and sparsification process is illustrated in Figure 1. Computationally, we may combine the above two goals into a unified objective for optimization:
max E A FEF
U [k) - [ = max Ez [< Z:
s.t. Z f(X), (1)
FEF
where the e norm Z o promotes the sparsity of the final token representations Z = f(X). We
call this objective "sparse rate reduction." "
White-box deep architecture as unrolled incremental optimization. Although easy to state, each term of the above objective can be computationally very challenging to optimize [55, 70J. Hence is natural to take approximation approach that realizes the global transformation f optimizing (I)
through a concatenation of multiple, say L, simple incremental and local operations f that push the representation distribution towards the desired parsimonious model distribution:
J: x
z
- ..... zt Z,
(2)
where : PRRD + Id is the pre-processing mapping that transforms input tokens æ, E RR to their token representations 27 E R2.
Each incremental forward mapping z = f' (z'), or a "layer", transforms the token distribution to optimize the above sparse rate reduction objective (1), conditioned on the distribution of its input tokens Z In contrast to other unrolled optimization approaches such as the ReduNet [55], we explicitly model the distribution of z' at each layer, say as a mixture of linear subspaces or sparsely generated from dictionary. The model parameters are learned from data (say via backward propagation with end-to-end training). This separation of forward "optimization" and backward "learning" clarifies the mathematical rolc of each layer an operator transforming the distribution of its input, whereas the input distribution is in turn modeled (and subsequently learned) by the parameters of the layer.
We show that we can derive these incremental, local operations through unrolled optimization perspective to achieve (1) through Sections 2.3 to 2.5. Once we decide on using an incremental approach to optimizing (1), there are a variety of possible choices to achieve the optimization. Given a model for Z say a mixture of subspaces U(. we opt for a two-step alternating minimization process with a strong conceptual basis: first in Section 2.3, we compress the tokens ze via gradient step to minimize the coding rate term minzz It ZZ U (K)) second, in Section 2.4, we spursify the compressed tokens, with a tuitbbl--relaxee proximal gradient step on the difference of the sparsity penalty and the expansion term, i.c., minz 2 Z - R( Z)]. Both actions are applied incrementally and repeatedly, as each fe in (2) is instantiated with these two steps.
2.2 Self-Attention via Denoising Tokens Towards Multiple Subspaces
There are many different ways to optimize the objective (1) incrementally. In this work, we propose arguably the most basic scheme. To help clarify the intuition behind our derivation and approximation, in this section (and Appendix 1) we study a largely idealized model which nevertheless captures the essence of nearly the whole process and particularly reveals the reason why self-attention-lict operators arise in many contexts. Assume that N = 1, and the single token æ is drawn i.i.d. from an unknown mixture of Gaussians (O, k supported on low-dimensional subspaces with K
k= k== k= and corrupted with additive Gaussian noise w ~ (O, 7), i.e.,
orthonormal bases U|K = (U+)
x 2 ow.
where z is distributed according to the mixture. Our goal is simply transform the distribution of the noisy token x to the mixture of -iwddimensional Gaussians z. Towards incremental construction of representation of for this model following (2), we reason inductively; if z is noisy token (3) at noise level of, it is natural to produce 2+1 by denoising at the level of. In the mean-square sense, the optimal estimate E[[z 2"), which has a variational characterization (e.g. [12]):
E -1 = arg min TE
[[][(z + oww =1(2]
2,10
Thha is. having the fewest nonzero entries. 'To simplify the notation, we will discuss the objective for one sample X at time with the understanding that we always mean optimize the expectation.

---

Setting =+1
E[z )) (4) thus characterizes the next stage of (2) in terms of an optimization
objective based on a local signal model for z!. Moreover, letting a + o (a:) denote the density of z', Tweedie's formula [13] allows us to express the optimal representation solving (4) in closed-form:
++1 = x + (o log (z').
(5)
Tweedie's formula expresses the optimal representation in terms of an additive correction (in general a nonlinear function of x") to the noisy observations by the gradient of the log- leglikkli of the distribution of the noisy observations, giving the optimal representation clear interpretation as an incremental perturbation to the current noisy distribution q This connection is well-known in the areas of estimation theory and inverse problems [1, 13, 14, 19, 20, 27, 42], and more recently has found powerful applications in the training of generative models for natural images [4, 15, 22, 43, 44]. Here, we can calculate a closed-form expression for this score function V log q", which, when combined with (5) and some technical assumptions", gives the following approximation (shown in Appendix A.I). Let & denote the Kronecker product; then we have
10
Uiz
1
2+1 Rs [UI..... Uk]
diag softmax
I
(6)
2(
LIUR>111
[
This operation resembles a self-attention layer in a standard transformer architecture with K heads, sequence length N 1, the gguery-key-value" constructs being replaced by a single linear projection Uf of the token z and the aggregation of head outputs (conventionally modeled by an MLP) done with the two leftmost matrices in (6). We thus derive the following useful interpretation, which we will exploit in the sequel: Gaussian denoising against a mixture of subspaces model leads to self-attention--yye layers in the transformation J. Given an initial sample æ following the model (3), we can repeatedly apply local transformations to the distribution with (6) in order to realize the incremental mapping f: æ + z in (2). 5 These insights will guide us in the design of our white-box transformer architecture in the upcoming subsections.
2.3 Self-Attention via Compressing Token Sets through Optimizing Rate Reduction
In the last subsection, we have seen that the multi-head attention a transformer resembles the score- matching operator that aims to transform a token 2f towards a mixture of subspaces (or degenerate Gaussians). Nevertheless, to carry out such an operation on any data, one needs to first learn or estimate, typically from finite samples, the parameters of the mixture of (degenerate) Gaussians, which is known to be a challenging task [6, 24]. This challenge is made even harder because in a typical learning setting, the given set of tokens are not i.i.d. samples from the mixture of subspaces. The joint distribution among these tokens can encode rich information about the data- example, co-occurrences between words or object parts in language and image data (esp.)-which. we should also learn. Thus, we should compress/ / denoise/ transform such a set of tokens together. To this end, we need measure of quality, i.e., compactness, for the resulting representation of the set of tokens. natural measure of the compactness of such a set of tokens is the (lossy) coding rate to encode them up to a certain precision > 0 [6. 46]. For a zero-mean Gaussian, this measure takes closed form. If we view the tokens in Z E RdxN as drawn from a single zero-mean Gaussian, an estimate of their (lossy) coding rate, subject to quantization precision > 0. is given in [6] as:
R(Z) ( logdet (I+ NZZZZZ 2 = logdet (I+ Ne
d
1
d
I+
Zz') Ne²
In practice, the data distribution typically multi-modal, say an image set consisting of many classes or a collection of image patches as in Figure 1. It is more appropriate to require that the set of tokens map to a mixture of, say K, subspaces (degenerate Gaussians) [55]. As before we denote the (to be learned) bases of these subspaces as UyK (U+) 1. where UR E Rd Although the joint distribution of the tokens Z is unknown, the desired marginal distribution of each token 2, is a
'Such as a being smaller than the nonzero eigenvalues of Σg and the normalization assumption क. det + -1/2 πy det(E,, + o² 1 for all i,j E [K]. where m. is the mixture proportion for the t Gaussian. 'This statement can be made mathematically rigorous by exploiting a deep connection between neural ODEs and diffusion models, following ideas in Song al. [44] and Chen et al. [72].
5

---

mixture of subspaces. So we may obtain an upper bound of the coding rate for the token set Z by projecting its tokens onto these subspaces and summing up the respective coding rates:
K
(Z: Σ R(
logdet (1+
P Ne (U(Z)
(8)
We would like to compress (or denoise) the set of tokens against these subspaces by minimizing the coding rate. The gradient of Rt (Z; (Z;Ukk) is
K ZUUZZ
(Z
p N
E
(I+ P Ne (22))(12
The above expression approximates the residual of each projected token U regressed by other tokens = [55]. But, differently from [55], not all tokens in Z are from the same subspace. Hence, to denoise each token with tokens from its own group, we can compute their similarity through an auto-correlation among the projected tokens as (U Z) (U; Z) and convert it to a distribution of membership with a softmax, namely softmax((U; 2 (U )). Then, as we show in Appendix A.2,
if we only use similar tokens regress and denoise each other, then a gradient step on the coding rate with learning rate can be naturally approximated as follows:
2+2 = Zt - xV z R
U(K) ~ (1 (1-*) K Ne
Z +
P
MSSA(Z |
UjK)), (10)
Ne
where MSSA is defined through an SSA operator as:
SSA(Z (k)) = (U Z) softmax 2) (U Z)),
E [K],
(11)
SSA Z U1)
MSSA Z U k))
N UU , Ux]
p
(12)
SSA Z Uk)
Here the SSA operator in (11) resembles the atlention operator in a typical transformer [28], except that here the linear operators of value, key, and query are all basis, i.c., V = K = Q U Hence, we name SSA(- (Ux) Rdx - RPXN N the Subspace Self-Attention (SSA) operator (more details and justification can be found (72) in Appendix A.2). Then, the whole MSSA operator in (12), formally defined as MSSA( [U]x]) Rdx + Rd and called the Multi-Head Subspace Self-Attention (MSSA) operator, aggregates the attention head outputs by averaging using model-dependent weights, similar in concept to the popular multi-head self-attention operator existing transformer networks. The overall gradient step (10) resembles the multi-head self-attention implemented with a skip connection in transformers.
be the same as the subspace
Notice that if we have N 1 tokens well as take aggressive gradient step (x 1) and tune the quantization errоr (e vp/N), the multi-head subspace self-attention operator in (12) becomes the ideal denoiser defined in (6), with the one minor difference that the aggregation of the heads is done by linear function here, while in (6) it is done by nonlinear coctuur-of-expeeee type function." This provides two very related interpretations of the multi head self-attention operator, as denoising and compression against mixture of low-dimessional subspaces.
2.4 MLP via Iterative
Theissaggg-Tressolddnnggg Algorithms (ISTA) for Sparse Coding
In the previous subsection, focused on how to compress set of tokens against set of (learned) low-dimensional subspaces. Optimizing the remaining terms in the sparse rate reduction objective (1), including the non-smooth term, serves to sparsify the compressed tokens, hence leading to more compact and structured (i.e., parsimonious) representation. From (1) and (7), this term is
d N zzZ Z
max [R(Z) - A| [b]] = min 2
Z| 10- - logdet
I+
(13)
2
"We note recent suggestion of Hinton [51] that it is sensible set the "value, key, and query" projection matrices a transformer to be equal. Our derivation this section confirms this mathematically. 'This suggests that we could also consider such mixture of expert type aggregation of the multiple attention heads. In this work, use linear aggregation, and leave evaluation of more variants for future work.

---

4
D
(D')*
LayerNorm
2/1
Activation
2
Sparse Coding Provimal Stop
(ISTA)
ISTA Block
2
Autocorrelation & Softrox
Add & JayerNorm
SSA Block
Head Subspace
Self-Artention
(MSSA)
(ncad1)
Aggrogate
4+122
z
SSA (hoodd
MSSA Block
Figure One layer of the CRATE architecture. The full architecture is simply concatenation of such layers, with some initial tokenizer and linal task snecispecific architecture (i.e., a classification head).
where R(Z) denotes the coding rate of the whole token set, as defined (7). In addition to
sparsification via the || term, the expansion term R Z)) in (13) promotes diversity and non- collapse the representation, a highly desirable property. However, prior work has struggled to realize this benefit on large-scale datasets due poor scalability of the gradient V z R( Z), which requires a matrix inverse [55].
To simplify things, we therefore take a different approach to trading off between representational diversity and sparsification: we posit (complete) incoherent orthogonal dictionary D E Rdxd and ask sparsify the intermediate iterates Z e++1/2 with respect D. That is, z Dz e+ where ze+1 more sparse. The dictionary D is global, i.e., used to sparsify all tokens simultaneously. By the incoherence assumption, we have D'D s Iai thus from (7) we have R R(Z s R(DZ*+1) = R(Z* 1112) Thus we approximately solve (13) with the following program:
z+1 arg min Zloo subject to 2+1 DZ.
(14)
The above sparse representation program is usually solved by relaxing it to unconstrained convex program, known as LASSO:
z arg min
\|| + 2
DZ
(15)
z
In our implementation, motivated by Sun et al. [33] and Zarka et al. [35], we also add non-negative constraint Z e+
Z a min
||
1 - DZ
(16)
Z20
which we then incrementally optimize by performing an unrolled proximal gradient descent step, known as an ISTA step [8], give the update:
7 +- ReLU (eLU + D DZ 1711) ISTA D).
(17)
In Appendix A.3, will show one can arrive at similar operator the above ISTA-like update for optimizing (13) by properly linearizing and approximating the rate term R(Z).
2.5 The Overall White-Box CRATE Architecture
By combining the above two steps:
1. (Sections 2.2 and 2.3) Local denoising and compression of tokens within a sample towards
mixture-of-see mixture-of-subspacc subspace structure, leading the multi-head subspace self-attention block - MMSSA;
7

---

2. (Section 2.4) Global compression and sparsification of token sets across all samples through sparse coding, leading to the sparsification block ISTA;
we can get the following rate-reduction-biie based transformer layer. illustrated in Figure 2,
2 = Z' MSSA(ZZ Z ((K])
Z = ISTA ISTA(Z D³).
(18)
Composing multiple such layers following the incremental construction of our representation in (2), we obtain white-box transformer architecture that transforms the data tokens towards compact and sparse union of incoherent subspaces.
This model has the parameters (Ukk))&= and (D)) which are learned from data via back- propagation. Notably, in each layer É, the learned Uik] retain their interpretation as incoherent bases for supporting subspaces for the mixture-of-Gaussian retains its interpretation as a sparsifying dictionary at layer É. We emphasize that the parameters UiK] and D are dependent on the layer L - that is, we learn different set of parameters at each layer. This is because at each layer learn an approximate local parametric model for the input data distribution, then use that learned model to construct the layer operators that transform the distribution. Our procedure of parameterizing the data distribution each layer distinguishes this work from previous works on unrolled optimization for neural networks such as the ReduNet [55]. Our interpretation clarifies the roles of the network forward pass (given local signal models each layer, denoiss/comprees//pprsii/ the input) and the backward pass (learn the local signal models from data via supervision).
model at layer e, and the learned De
We note that in this work, at each stage of our construction, we have chosen arguably the simplest possible construction to use. We can substitute each part of this construction, so long as the new part maintains the same conceptual role, and obtain another white-box architecture. Nevertheless, our such-constructed architecture, called CRATE (i.e., Coding RAte TransformFr), connects to existing transformer models, obtains competitive results on real-world datasets, and is fully mathematically interpretable.
3 Experiments
In this section, we conduct experiments to study the performance of our proposed white-box trans- former CRATE on real-world datasets and tasks. As the analysis in Section 2 suggests, either the compression or the sparsification step can be achieved through various alternative design choices or strategies. CRATE arguably adopts the most basic choices and so our goal with the experiments not simply to compete with other heavily engineered transformers while using such a rudimentary design. Rather, our goals are twofold. First, unlike any empirically designed black-box networks that are usually evaluated only on end-to-end performance, the white-box design of our network allows us to look inside the deep architecture and verify layers of the learned network indeed perform their design objective- -sy performing incremental optimization for the objective (1). Second, despite their simplicity, our experiments will actually reveal the vast practical potential of our so-dcrived CRATE architectures since, as we will show, they already achieve very strong performance on large-scale real-world datasets and tasks. In the remainder of this section we highlight a selection of results; additional experimental details and results can be found in Appendix B.
Model architecture. We implement the architecture that is described in Section 2.5, with minor modifications that are described in Appendix B. We consider different model sizes of CRATE by varying the token dimension d, number of heads K, and the number of layers L. We consider four model sizes in this work: CRATE-Timy: CRATE-Small CRATE- Base, and - A PyTorch- style pseudocode can be found in Appendix B. which contains more implementation details. For training using supervised classification, we first take the CLS token 3 = 21,6
L+1 of for each sample,
then apply linear layer; the output of this linear layer u = Wz, is used input to the standard cross-entropy loss. The overall loss averages over all samples 5 E [B].
Datasets and optimization. We mainly consider ImageNet- K [9] as the testbed for our architecture. Specifically, we apply the Lion optimizer [73] to train CRATE models with different model sizes. Meanwhile, we also evaluate the transfer learning performance of CRATE: by considering the models trained on Image 1K as pre-trained models, we fine-tune CRATE on several commonly used downstream datasets 10/100 Oxford Flowers, Oxford-IIT-Pets). More details about the training and datasets can be found in Appendix B.I.
8

---

Measure ooding rate across layers
Measure output sparsity acrees layers
R1222) ISSA block
Sparsity
train
train val
68
Layer Index
Layer Index -
Figure 3: Left: The compression term (Z(+11/2 of the MSSA outputs at different layers. Right: the sparsity the ISTA output block, 1 lo/(d N). at differrnt layers. (Model: CRATE- CRATE-Small)
Measure coding rate across layors
Measure oupput sparsity across bayers
1000
(ISTA block)
irit.
epoch
epoch 1
cpach
opach
epoch 150
spach 150
Layer Index
Layer index
Figure 4: The compression term (Z) (left) and sparsification term 2 lo/ N) (right) across models trained with different numbers of epochs. (Model: CRATE-Basee Basee..
3.1 In-depth Layer-wise Analysis of CRATE
Do layers of CRATE achieve their design goals? As described in Section 2.3 and Section 2.4, the MSSA block is designed optimize the compression term Rt and the ISTA block sparsify the token representations (corresponding to the sparsification term Z llo). To understand whether CRATE
indeed optimizes these terms, for each layer f, we measure (i) the compression term R ( 1221
on the MSSA block outputs Z 1+1/2 and (ii) sparsity 1 on the ISTA block outputs z Specifically, evaluate these terms by using training/validationn samples from ImageNet-IKK IK. Both terms are evaluated at the per-sample level and averaged over B = 10 samples.
Figure 3 shows the plots of these two key measures at all layers for the learned CRATE-smal model. We find that as the layer index é increases, both the compression and the sparsification terms improve in most cases. The increase in the sparsity measure of the last layer is caused by the extra linear layer for classification." 8 These results suggest that CRATE aligns well with the original design goals: once learned, it essentially learns to gradually compress and sparsity the representations through its layers. In addition, we also measure the compression and sparsification terms on CRATE models with different model sizes well intermediate model checkpoints and the results are shown by plots in Figure 5 of Appendix B.2. The observations are very consistent across all different model sizes-bote the compression and sparsification terms improve in most scenarios. Models with more layers tend to optimize the objectives more effectively, confirming our understanding of each layer's roles.
To the effect of learning, we present the evaluations CRATE-Small trained with different number of epochs in Figure 4. When the model is not trained enough (e.g. untrained), the architecture does not optimize the objectives effectively. However, during leannigg better subspaces UjK] and dictionaries D the designed blocks start to optimize the objectives much more effectively.
Visualizing layer-wise token representations. To gain better understanding of the token represen- tations of CRATE, we visualize the output of each ISTA block at layer in Figure 6 of Appendix B.2. Specifically, visualize the z via heatmap plots. We observe that the output Z becomes more sparse as the layer increases. Moreover, besides the sparsity, we also find that e++ becomes
Note that the learned sparse (tokens) features need to be mixed in the last layer for predicting the class. The phenomenon of increase in the sparsity measure at the last layer suggests that each class of objects may be associated with number of features, and some of these features are likely to be shared across different classes.

---

Table I: Top accuracy of CRATE various datasets with different model scales when pre-trained on ImageNect. For ImageNe//mageNNeRReal, directly evaluate the top-1 accuracy. For other datasets, models that are pre-trained ImageNet as initialization and the evaluate the transfer learning performance fine-tuning.
Datasets
CRATE-T
CRATE-S
CRATE-B
CRATE-L
VIF-T
VIT-S
# parameters
6.09M
13.12M
22.80M
77.64M
5.72M
22.05M
ImageNet
66.7
69.2
70.8
71.3
71.5
72.4
ImageNet ReaL
74.0
76.5
77.4
78.3
78.4
CIFARIO
95.5
96.0
96.8
97.2
96.6
97.2
CIFARIO0
78.9
81.0
82.7
83.6
81.8
83.2
Oxford Flowers- 102
84.6
87.1
88.7
88.3
85.1
88.5
Oxford-IIT-Pett
81.4
84.9
85.3
87.4
88.5
88.6
more structured (i.e., low-rank), which indicates that the set of token representations become closer to linear subspaces, confirming mental picture of the geometry of each layer (as in Figure 1).
Visualizing layer-wise subspaces in multi-head self-attention. We now visualize the U(K] ma- trices used in the MSSA block. In Section 2.3, assumed that [KK were incoherent to capture different "views" of the set of tokens. In Fig. 7 of Appendix B.2, first normalize the columns in each U then visualize the [U U [U s..... Ukl E RoK The (i. (i.j)-th --th block in each sub-figure corresponds to Uf for i,j [K] particular layer L. We find that the learned [K] are approximately incoherent, which aligns well with our assumptions. One interesting observation is that the U becomes more incoherent when the layer index is larger, which suggests that the token representations more separable. This mirrors the situation in other popular deep networks [58].
3.2 Evalutions of CRATE on Large Real-World World Datasets and Tasks
We now study the empirical performance of the proposed networks by measuring their top-1 accuracy on I K well as transfer learning performance several widely used downstream datasets. We summarize the results in Table 1. As our designed architecture leverages parameter sharing in both the attention block (MSSA) and the MLP block (ISTA), our CRATE-Base model (22.08 million) has a similar number of parameters to the ViT-Small (22.05 million).
From Table 1, we find that with a similar number of model parameters, our proposed network achieves similar ImageNet-I and transfer learning performance as ViT, despite the simplicity and interpretability of our design. Moreover, with the same set of training hyperparameters, we observe promising scaling behavior in CRATE we consistently improve the performance by scaling up the model size. For comparison, directly scaling ViT on ImageNet-1K K does not always lead to consistent performance improvement measured by top-1 accuracy [40]. To summarize, we achieve promising performance on real-world large-scale datasets by directly implementing our principled architecture.
4 Conclusion
In this paper, we propose a new theoretical framework that allows us to derive deep transformer- like network architectures incremental optimization schemes to learn compressed and sparse representation the input data (or token sets). The so derived and learned deep architectures are not only fully mathematically interpretable, but also consistent objective. Despite being arguably the simplest among all possible designs, these networks already demonstrate performance on large-scale real-world datasets and tasks close to seasoned transformers. We believe this work truly helps bridge the gap between theory and practice of deep neural networks as well help unify seemingly separate approaches to learning and representing data distributions. Probably more importantly for practitioners, our framework provides theoretical guidelines to design and justify new, potentially more powerful, deep architectures for representation learning.
layer-by-layer level with their design
10

---

Acknowledgements
We thank the anonymous reviewers for their helpful comments. Yaodong Yu would like to thank Kwan Ho Ryan Chan for the valuable discussions we had regarding visualizing tokens in vision transformers. Yaodong Yu acknowledges support from the joint Simons Foundation-NST
DMS grant #2031899. Yi Ma acknowledges support from ONR grant N00014-22- 22-1-21022 and the joint Simons Foundation-NsT. DMS grant #203 1899. This work was partially supported by NSF 1704458, the Northrop Grumman Mission Systems Research in Applications for Learning Machines (REALM) initiative, NIH NIA AGO19966, and ARO MURI
References
[1] Charles M Stein. "Estimation of the Mean of Multivariate Normal Distribution". The Annals of Statistics 9.6 (Nov. 1981), pp. 1135-1151. 5.
[2] Bruno A Olshausen and David J Field. "Sparse coding with an overcomplete basis set: A strategy employed by VI?" Vision research 37.23 (1997), pp. 33113331323 1-3325. 2.
[3] David L Donoho and Carrie Grimes. "Image Manifolds which are Isometric to Euclidean Space" Journal of mathematical imaging and vision 23.1 (July 2005), pp. 5-24. 2.
[4] Aapo Hyvärinen. "Estimation of Non-Normalized Statistical Models by Score Matching". Journal of machine learning research: JMLR 6.24 (2005), pp. 695-709. 5.
[5] Michael B Wakin, David L Donoho, Hyeokho Choi, and Richard G Baraniuk. "The multiscale structure of non-differentiable image manifolds". Wavelets XI. Vol. 5914. SPIE. 2005, pp. 413- 429, 2.
[6] Yi Ma, Harm Derksen, Wei Hong, and John Wright. "Segmentation of multivariate mixed data via lossy data coding and compression". PAMI (2007). 2, 3, 5.
[7] Maria-Elena Nilsback and Andrew Zisserman. "Automated flower classification over large number of classes" 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing. IEEE 2008, pp. 722-729 26.
[8] Amir Beck and Mare Teboulle. "A fast iterative shrinkage- taaikkaggg-theehhlnnngg algorithm for linear inverse problems". SIAM journal on imaging sciences 2.1 (2009). pp. 183-202. 7.
[9] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. "Imagenet: A large- scale hierarchical image database". 2009 IEEE conference on computer vision and pattern recognition. leee. 2009, pp. 248-255. 8, 35.
[10] Alex Krizhevsky, Geoffrey Hinton, et al. "Learning multiple layers of features from tiny images" (2009). 26.
[11] Karol Gregor and Yann LeCun. "Learning fast approximations of sparse coding". Proceedings the 27th International Conference on International Conference on Machine Learning. Omnipress. 2010, pp. 399-406. 2.
[12] László Györfi, Michael Kohler, Adam Krzyzak, and Harro Walk. A Distribution-Fred Theory of Nonparametric Regression. Springer New York, Dec. 2010. 4.
[13] Bradley Efron. "Tweedie's s Formula and Selection Bias". Journal of the American Statistical Association 106.496 (201 1), pp. 1602-1614. 2, 16.
[14] Martin Raphan and Eero P Simoncelli. "Least squares estimation without priors or supervision". Neural computation 23.2 (Feb. 2011), pp. 374-420. 5.
[15] Pascal Vincent. "A connection between score matching and denoising autoencoders". Neural computation 23.7 (July 2011), pp. 1661-1674.
[16] Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. "Cats and d 2012 IEEE conference on computer vision and pattern recognition. IEEE. 2012, pp. 3498-3505. 26.
[17] Daniel Spielman, Huan Wang, and John Wright. "Exact Recovery of Sparsely- Dictio- naries" (June 2012). arXiv: 2206.5883
LG]. 2.
[18] Joan Bruna and Stéphane Mallat. "Invariant scattering convolution networks". IEEE transac- tions on pattern analysis and machine intelligence 35.8 (Aug. 2013), pp. 1872-1886.
[19] Peyman Milanfar. "A Tour of Modern Image Filtering: New Insights and Methods, Both Practical and Theoretical". IEEE Signal Processing Magazine 30.1 (Jan. 2013), pp. 106-128.
5.

---

[20] Singanallur V Venkatakrishnan, Charles A Bouman, and Brendt Wohlberg. "Plug Badu--nd--o and--ad-PPaa Play pri- ors for model based reconstruction". 2013 IEEE Global Conference on Signal and Information Processing. Dec. 2013, pp. 945-948. 5.
[21] Rémi Gribonval, Rodolphe Jenatton, and Francis Bach. "Sparse and spurious: dictionary learning with noise and outliers" (July 2014). arXiv: 1407. 5155 [cs LG]
[22] Jascha Sohh-Dickstein, Eric A Weiss, Niru Maheswaranathan, and Surya Ganguli. "Deep Unsupervised Learning using Nonequilibrium Thermodynamics" (Mar. 2015). arXiv: 1503.
03585 [ca. [cs.LG] 2.
[23] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. "Deep Residual Learning for Image Recognition" " 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). June 2016. pp. 770-778. 3,.
[24] René Vidal, Yi Ma, and Shankar Sastry. Generalized Principal Component Analysis. Springer Verlag, 2016. 5.
[25] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. "Mask R-CNN" (Mar. 2017). arXiv: 06870 [cs. cv].
[26] Ilya Loshchilov and Frank Hutter. "Decoupled weight decay regularization". arXiv preprint
arXiv:1711.055 (2017). 26.
[27] Yaniv Romano, Michael Elad, and Peyman Milanfar. "The Little Engine That Could: Regular- ization by Denoising (RED)". SIAM journal on imaging sciences 10.4 (Jan. 2017), pp. 1804- 1844. 5.
[28]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. Lukasz Kaiser, and Illia Polosukhin. "Attention is all you need". Advances in neural informa- tion processing systems 30 (2017). 1-3,6 6.
[29] Yubei Chen, Dylan Paiton, and Bruno Olshausen. "The sparse manifold transform". Advances in neural information processing systems 31 (2018). 2.
[30] Jacob Devlin, Ming-We Wei Chang. Kenton Lee, and Kristina Toutanova. "Bert: Pre-training of deep bidirectional transformers for language understanding". arKiv preprint (2018). I.
[31] Tero Karras, Samuli Laine, and Timo Aila. "A Style-Based Generator Architecture for Genera- tive Adversarial Networks" (Dec. 2018). arXiv: 1812. 049488 [cs. .N]. 1.
[32] Vardan Papyan, Yaniv Romano, Jeremias Sulam, and Michael Elad. "Theoretical Foundations of Deep Learning via Sparse Representations: A Multilayer Sparse Model and Its Connection to Convolutional Neural Networks". IEEE Signal Processing Magazine 35.4 (July 2018),
pp. 72-89. 2.
[33] Xiaoxia Sun, Nasser M Nasrabadi, and Trac D Tran. "Supervised deep sparse coding networks". 2018 25th IEEE International Conference on Image Processing (ICIP). IEEE. 2018, pp. 346- 350. 7.
[34] Yang Song and Stefano Ermon. "Generative Modeling by Estimating Gradients of the Data Distribution" (July 2019). arxiv: 1907. 05600 [cs LG). 2.
[35] John Zarka, Louis Thiry, Tomás Angles, and Stéphane Mallat. "Deep network classification
by scattering and homotopy dictionary learning". arXiv preprini arXiiv/990.000 (2019).
[36] Lucas Beyer, Olivier Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. "Are we done with imagenet?" arXiv preprint 20xiv:2006607153 (2020). 26.
[37] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. "Language models are fhw-shot learners". Advances in neural information processing systems 33 (2020), pp. 1877-1901. 1.
[38] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. "End-to-Be to-dnt Object Detection with Transformers" (May 2020), arXiv: 1007. [cs.CV]. 1.
[39] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. "A Simple Frame- work for Contrastive Learning of Visual Representations" Proceedings of the 37th Interna- tional Conference on Machine Learning. Ed. by Hal Daumé lii and Aarti Singh. Vol. 119. Proceedings of Machine Learning Research. PMLR, 2020, pp. 1597-1607. 1.
12

---

[40] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly. et al. "An image is worth 16x 16 words: Transformers for image recognition at scale". arXiv preprint arXiv:2010. 1929 (2020). 1, 3, 10, 36.
[41] Jonathan Ho. Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models". vances in Neural Information Processing Systems 33 (2020), pp. 6840-6851. 2.
[42] Zahra Kadkhodaic and Eero P Simoncelli. "Solving Linear Inverse Problems Using the Prior Implicit in a Denoiser" (July 2020). arXiv: 2007. 13640 ccs.cV]. 5.
[43] Jiaming Song. Chenlin Meng, and Stefano Ermon. "Denoising Diffusion Implicit Models"
(Oct. 2020). arXiv: 2010. 02502 (Cs.LG]. 2, 5.
[44] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon,
and Ben Poole. "Score-Based" Generative Modeling through Stochastic Differential Equations"
(Nov. 2020). arviv: 2011 13456 LGs.LG]. 5.
[45] Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. "What makes for good views for contrastive learning?" Advances in neural information pro-
cessing systems 33 (2020), pp. 6827-6839. 2.
[46] Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. "Learning Diverse and Discriminative Representations via the Principle of Maximal Coding Rate Reduction". Advances in Neural Information Processing Systems 33 (2020), pp. 9122-94222999 9434. 2, 3, 5, 19, 24, 35.
[47] Yuexiang Zhai, Zitong Yang, Zhenyu Liao, John Wright, and Yi Ma. "Complete dictionary learning via 4-norm maximization over the orthogonal group". The Journal of Machine Learning Research 21.1 (2020), pp. 6622-6689. 2.
[48] Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. "Vivit: A video vision transformer". Proceedings of the IEEE/CVF international conference computer vision. 2021, pp. 6836-6844 6846.686.11 1.
[49] Florentin Guth, John Zarka, and Stephane Mallat. "Phase collapse in neural networks". arXivv preprint arXiv:2110.0522 (2021). 35.
[50] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. "Masked Autoencoders Are Scalable Vision Learners" (Nov. 2021). arXiv: 2111 06377 [vs.CV]. 1. [51] Geoffrey Hinton. How to represent part-whole hierarchies in a neural network. 2021. arXiv:
2102. 12627 [cs.CV].6 6.
[52] Alec Radford, Jong Wook Kim. Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. "Learning Transferable Visual Models From Natural Language Supervision". Pro- ceedings of the 38th International Conference on Machine Learning. Ed. by Marina Meila and Tong Zhang. Vol. 139, Proceedings of Machine Learning Research. PMLR, 2021, pp. 8748- 8763. 1.
[53] Bahareh Tolooshams and Demba Ba. "Stable and Interpretable Unrolled Dictionary Learning.. arXiv preprint wwXiv:2/06.0005: (2021). 2.
[54] Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung. Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. "MLP-Mixer: An all-MLP Architecture for Vision" (May 2021). arXiv: 2105. 01601 [cs.CV]. 23.
[55] Kwan Ho Ryan Chan, Yaodong Yu, Chong You, Haozhi Qi, John Wright, and Ma. NReduNet: A White-box Deep Network from the Principle of Maximizing Rate Reduction" Journal of Machine Learning Research 23. 14 (2022), pp. 1-103. 2-8, 19, 20.
[56] Hongrui Chen, Holden Lee, and Jianfeng Lu. "Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions", arxiv preprint arXiv:2211.011 (2022). 2.
[57] Yuan Gong, Andrew Rouditchenko, Alexander H Liu, David Harwath, Leonid Karlinsky, Hilde Kuehne, and James R Glass. "Contrastive audio-visual masked autoencoder". The Eleventh International Conference on Learning Representations. 2022. 1.
[58] Hangfeng He and Weijie Su. "A law of data separation in deep learning". arXiv preprint arXiv:2210. 17020 (2022). 10.
13

---

[59] Geoffrey Hinton. The Cerward-Forwate Algorithm: Some Preliminary Investigations. 2022.
arXiv: 2212. 13345 [Gs. 2.
[60] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. "Elucidating the design space of diffusion- based generative models". arXiv preprint arXiv:2206. 00364 (2022). 2, 16.
[61] Frederic Koehler, Alexander Heckett, and Andrej Risteski. "Statistical Efficiency of Score
Matching: The View from Isoperimetry" (Oct. 2022). arXiv: 2210.00726 00726 [cs LL]. 2.
[62] Yi Ma. Doris Tsao, and Hleung-Yeung Shum. "On the principles of parsimony and self- consistency for the emergence intelligence", Frontiers of Information Technology & Elec- tronic Engineering 23.9 (2022), pp. 1298-1323. 1,
[63] Druv Pai, Michael Psenka, Chih- Yuan Chiu, Manxi Wu, Edgar Dobriban, and Yi Ma. "Pursuit of a discriminative representation for multiple subspaces via sequential games". arXiv preprint arXiv:2206.0913 (2022). 2.
[64] Mary Phuong and Marcus Hutter. "Formal algorithms for transformers". arXiv preprint arXiv:2207.0923. (2022). 21.
[65] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. "High-resolution image synthesis with latent diffusion models". Proceedings of the IEFE/CVF Conference on Computer Vision and Pattern Recognition. 2022, pp. 10684-10695. 2.
[66] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David Fleet, and Mohammad Norouzi. "Photorealistic Text-to- Image Diffusion Models with Deep Language Understanding" (May 2022). arXiv: 2205 11487 [cs.Cv]. 1.
[67] Asher Trockman, Devin Willmott, and J Zico Kolter. "Understanding the Covariance Structure of Convolutional Filters" (Oct. 2022). arXiv: 2210. 03651 [cs.CV]. CV]. 23.
[68] Rene Vidal. Attention: Self-Expressiin All You Need. Unpublished; available: https
// openreviee appenreview.net/forrreee mett/forummidddMijBBo Micd 2022. 2.
[69]
Haoqing Wang, Xun Guo, Zhi-Hong Deng, and Yan Lu. "Rethinking minimal sufficient representation in contrastive learning". Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. pp. 16041-16050. 2.
[70] John Wright and Yi Ma. High-Dimenssinal Data Analysis with Low-Dimensional Models: Principles, Computation, and Applications. Cambridge University Press, 2022. 4, 22, 23.
[71] Yongyi Yang, Zengfeng Huang, and David P Wipf. "Transformers from an Optimization Perspective". Advances in Neural Information Processing Systems. Ed. by S Koyejo, Mo- hamed, A Agarwal, D> Belgrave, K Cho, and A Oh. Vol. 35. Curran Associates, Inc., 2022, pp. 36958-36971. 2.
[72] Sitan Chen, Giannis Daras, and Alexandros G Dimakis. netiooratioo-eeggaatttiii Beyond Linear Diffusions: A Asymptoott Analysis For DDIM-Type Samplers" (Mar. 2023). arXiv: 2303. 03384 [cs LG).
[73] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hicu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, et al. "Symbolic discovery of optimization algorithms". arXiv preprint arXiv:2332.0666
(2023). 8, 26.
[74] Brent De Weerdt, Yonina C Eldar, and Nikos Deligiannis. "Designing Transformer Networks for Sparse Recovery of Sequential Data Using Deep Unfolding"... ICASSP 2023 -023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). June 2023,
pp. 1-5.2.
[75] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, al. "Scaling vision transformers to 22 billion parameters". arXiv preprint arXiv:2302.05442 (2023). 1.
[76] Benjamin Hoover, Yuchen Liang, Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Horn Chau, Mohammed Zaki, and Dmitry Krotov. "Energy Transformer" (Feb. 2023). arXiv: 2302. 07253 [CG. 2.
[77] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, Piotr Dollár, and Ross Girshick. "Segment Anything" (Apr. 2023), arXiv: 2304. 02643 [cs.CV) 1.
14

---

[78] Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. "A Theoretical Understanding of shallow Vision Transformers: Learning. Generalization, and Sample Complexity". arXiv preprint arXiv: 06015 (2023). 2.
[79] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, and Sanjiv Kumar. "The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers", The Eleventh International Conference on Learning Representations. 2023. 23.
[80] Ravid Shwartz-Ziv and Yann LeCun. "To Compress or Not to
Learning and Information Theory: A Review". arXiv preprint arXiv:2304. 09355 (2023). 2.
[81] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. "Consistency models". arXiv preprint 2w3Xvv2303.00 (2023). 2.
15

---

Appendix
A Technical Details from Section 2
A.1 Companion to Section 2.2
We first wish to re-iterate the core contributions of our approach in Section 2.2 slightly more technical level. Connections between denoising and score matching well-understood [60], and computing the optimal denoising function (i.e., the conditional expectation) against a mixture-of- Gaussians model is rather simple computation giving existing tools such Tweedie's 's formula [13]. These are not our main contributions. Instead, the main contributions Section 2.2 are two-fold:
First, we demonstrate a mechanism to learn representations via denoising within idealized mixture of Gaussian data model for single token (i.e., with sequence length N 1).
Second, illustrate the similarities between such-derived representation learning scheme and existing self-attention layers within the transformer (with sequence length 1). thus demonstrating interpretation of the self-attention layer generalized mechanism to denoise against a mixture-of SaaereoofGGussssss Gaussian Gaussian-marginal model for a set of tokens.
Now we produce the proofs alluded to in Section 2.2, which mostly form the technical aspects of the first listed contribution. To simplify the proofs, we use the following notation correspondences: + z¹, + zt1,, and + of.
Proposition 1. Let u1,..., uk R&d be independent and have distribution un ~ (N Ex) for Eg 2 O, and let z take value un with probability Tx > O. Let w ~ N(0, Ia) be independent of z. Let ow. Let x q(æ) be the density of x. We define
M
kE + d² Ia) 11)-1/2
(19)
and assume that n. (e) = Tyy det( for all S S S K. Then we have
V log g(x))
(20)
| x 12
Mi
=- [M ... Mr] diag softmax
&
Ia
(21)
2
1|M
M
where denotes the Kronecker product, i.e., the block matrix defined by
An
A B
A & B
(22)
Am B
Amn B
Proof. Let u be the multinomial random variable such that z = so that u has probability mass function æ. Then by the law of total probability, we have
V (xx) V log Σ q(x
K)TT
(23)
k=1
k= Th V V k)
(24)
Σk k q(x K)TT
where q( kt) is the conditional density of given the event {u k}. To compute this quantity, note that conditional on the value of u, we have
Zu ow (O, E.. La)))
(25)
Thus we have
1 +=I
q(x: k)
exp
- -
(26)
(2 dEl I2)a
This gives
Vag(2 / -990 kt) (Ex o² o²La) -a'a)_aa x.
(27)
16

---

Putting this all together, we get
V x log ((:)
(28)
Σ 1(x 1TTT (Eis + 02
(29)
Σ 9(0 1T
k=1 The det Eg + 0² I1) 1/2 exp 2xpp-220'' (Ex + (aaa)aaa (En: + 02 -'
(30)
Ex= del Σ 0² 1 exp( <>x + =2 la))
Now define M2 = EEx + Ia)
1/2. -1/2, With this notation, we have
V logg( --
Ek=1 Th (M exp(-3a"MMMMaa) M
(31)
E Th det M&) exp( 2 M
K 2k k- ThH MM) exp( - a M
(32)
Σ T det (M) exp( 2 M²
Given our assumption that each det(M²)) is the same, we have
VR log ((x))
(33)
Σk= Th det M&) exp | M
(34)
Σ T det (M+) exp M
1= exp( M x 13) M
(35)
Σ exp(
M|
1
eg softmax
M.M
(36)
2
M
x
2
Mf
M Mx]
diag
softmax
-
@ La
(37)
M
Now we provide a final justification for the result cited in Section 2.2.
Approximation the setting of Proposition diagonalize Σn U where Ug Rdxp orthogonal and Ag TRxx diagonal." 9 Then we have the approximation
1
U
z
x
U1.....
Uk]
diag
softmax
Ip
:
(38)
20
U
Proof. We have
Va logg(
Σ k=1
C softmax
McM
(39)
[| 2
|| M M 20
Σ
ef softmax
M,MEa
(40)
2
[| M
9This assumption can be easily relaxed to As 2 0 for all & but requires some more notation to handle, and the form of the solution does not change. Thus we handle the case where all matrices full rank for simplicity.
17

---

||
o x 2
1
E
softmax
M,Mja.
(41)
20
||0:
|1o M 1
Now define Ps La- - Ma, and let Up E [dd((dp) p) be an orthogonal complement Ug. Then we have
Ps = - M
(42)
= Ld
Ex + 0² I'Ia)
-1/2
(43)
I
U Uk!
Aa 0
(U )r
U
+
+o01d
-1/2
(44)
La
U UH]
A
+ 0
o p
(UF))
U2
-1/2
(45)
PIa-p
P1
U U
(AAA + Ip))
1/2
0
U
(46)
(0²)-1/2 1/2
la
T1
U
U
(a
A
+
Ip)
1/2
Up
(47)
Ia-p_ pt
(U, )r
I
(a
A
I,))
11/
0
Up
U
U
(48)
0
0
(U r
?8 0
0
U
2 U
U
(49)
(U
= --
(50)
Thus P& approximately projection when is small. Under this algebraic relation, we have
V log ((x)
(51)
||0|||
||0 MMMiaa 2
K
1
e softmax
M,
(52)
20
||
||o M i [
||
[T p1)'
K
softmax
(Is P
PA)
(53)
20
k=1
|| x|
11Aa - Pk) x0
K
ex softmax
(Ia - Pk) (A)''
(54)
20
||P
K
E softmax
:
(Iu Pk)
(55)
20
|| [|3
K
1
softmax
c softmax
20
20
k=1
||
12
(56)
1
K
L +
ex softmax
P/
(57)
20²
km l
||
1
10
K
1
E +
ex
softmax
U
(58)
02
202
k=
| x |
18

---

1
U
1 1 o2 x +
नt [UI,... Ukl
diag
softmax
:
1p
:
(59)
2x
|
U
Plugging this into Tweedie's formula, have
|0
Uf
1
B 2 x 22
Ukl
diag
softmax
② 1p
(60)
2x²
U
Remark 3. Although Approximation 2 is stated as approximation rather than proposition, believe it should possible without too much extra work to convert into a statement asymptotic equivalence - 0 (in particular, holding for below the smallest (nonzero) eigenvalue of any Σg. Most approximations taken in the derivation Approximation can immediately be turned into asymptotic claims; the only slightly delicate point is treating the softmax, which can be accomplished using standard "high temperature" convergence behavior of the softmax function (in particular, as in our expressions, the softmax concentrates on the "best head").
A.2 Companion to Section 2.3
We again wish to re-iterate the core contribution of our approach in Section 2.3. The application of a compression perspective representation learning has been discussed before, for example in the line of maximal coding rate reduction works [46]. In Section 2.3, we provide the following contributions and developments to this perspective:
We propose generalized coding rate function (:; U(K) which measures the coding rate with respect to a set of subspaces UKK] as opposed set of classes (as in [46, 55]), making the underlying formulation unsupervised.
We then show how if we adopt the framework of alternating minimization of the sparse rate reduction objective, then unrolling the first alternating step - gradient descent on this coding rate objective nearly exactly recovers the common multi-head attention mechanism found in transformer networks (except that the query/key/value operators are all the same operation Uj now, which we interpret projection onto single subspace).
In the process of the second contribution, and in the following proofs, we make some simple approximations and technical assumptions. The validity of these assumptions may be explored, and the approximations refined, altogether providing a more complex (and possibly more performant) resulting self-attention like operator. For the sake of technical clarity and simplicity in this work, make perhaps the simplest possible choices. As a result, we do not claim that our network is optimally designed, but rather that the principles unrolled optimization) can provide the backbone for far superior and more interpretable network architectures in the future on sundry tasks. As it is, with design, we still obtain meaningful conceptual results and very solid empirical performance.
develop in this work (compression, denoising, sparsification,
straightforward, simple, and interpretable
We now give the derivation of the approximation alluded to in Section 2.3.
Approximation Let Z R have unit-norm columns, and UlK (UI...., ...,,K)) such that each U* ddxp an orthogonal matrix, the (=)) are incoherent, and the columns of Z approximately lie on U Span(Uk). Let NZz. Let 0. Then
-- Z UjK) s (11x,))2 kry MSSA ((44)
(61)
where in Section 2.3 we have
SSA (Z))) (U, Z) softmax
(U; Z)),
(62)
SSA(Z(U1)
MSSA (SSA(Z((KK) IU...... Ukl
(63)
SSA( (Uk)
19

---

where softmax(.) is the softmax operator (applied to each column of an input matrix), i..,
softmax(v)
(64)
E en
6
softmax((v),..... vx() = (oottmax(Ui)...oo softmax(Ux])
(6.5)
Proof. According to (9), the gradient Zz U(K) (x))
is
K
(2; - -000011111
(66)
k=1
Notice that according to [55], the gradient is precisely the residual of a ridge regression for each (projected) token U zi using other projected tokens U 2, as the regressors, hence being the residual of an auto-regression.
have (I + (U, (U; Z)) be expensive. Hence for computational efficiency, may approxi- mate with the first order term of its von Neumann expansion:
However,
in the work of ReduNet [55]. computing the inverse
V (2
E TAU;Z(1+N1)) (1
(67)
R 7
U U1U;
(68)
k=1
U (,00;3--U;220000000000 Σ k=
(69)
Notice that the term ZU is the auto-correlation among the projected tokens. As the tokens Z may be from different subspaces, we would prefer to use only tokens that belong to the same subspace to regress and compress themselves. Hence we may convert the above correlation term into a subspace-membershe indicator with a softmax operation, whence (69) becomes
V (2
Σ Un -U;
(70)
K
Σ
U - ? E Uk (U SUK(URz softmax (U1
(71)
Then, we can rewrite the above approximation the gradient of R as:
K
K
Vz (2 =
Σ
U -
Σ
Un (U; 2 (0t2))
(72)
k=1
k=
K
U -72
U₁ SSA(Z
Ux)
(73)
k=1
SSA(Z U1)
U Z Z- --,2
Unl
(74)
SSA(Z Ux)
SSA U!)
2 22 72 U
Unl
(75)
SSA(Z
Uk)
20

---

Thus the gradient descent step with learning rate >> 0 ggives
SSA(Z :
- xV ( ~ (1-m7) +m7² 00 U Uk]
(76)
SSA(ZIUR)
A.3
Companion to Section 2.4
We again wish to re-iterate the core contribution of our approach in Section 2.4.
Within the framework of alternating minimization of the sparse rate reduction objective, show that the second alternating step gradient descent on the overall coding rate plus a sparse regularization term has heuristic connections to particular LASSO optimization.
We show that the unrolling of the proximal gradient step to solve this LASSO optimization resembles the MLP which immediately follows the self-attention layer within transformer blocks.
In the main text, our connection between the second step of the alternating minimization and the LASSO optimization was high-level and heuristic. In some sense, the choice pose the minimization step as a LASSO was a simple, reliable, and interpretable choice which works well in practice, but nonetheless not backed up by rigorous theoretical justification. the following subsection, we provide mathematical justification for reformulation of the minimization step using majorization- minimization framework. We further show that the associated unrolled optimization step bears a strong resemblance to the ISTA step. This confirms our earlier discussion possible choice in designing CRATE, but by more rigorous derivation we can uncover alternative operators which nonetheless have the same conceptual function and may perform better in practice.
took the simplest
Assumptions. In this section, we present rigorous optimization analysis of an incremental minimization approach to the objective (13). We will show that under two simplifying assumptions, namely
The columns of 2++1/2 are normalized, the sense that diag( Z(+1/2) = 1:
2. We have 2 N, || and the columns Z 1/21/2 are orthogonal, so that (Z#+1/2)2 2t 1/2
I. I.
the approach leads to an update iteration that is equal to a slightly simplified version of the ISTA block (17). We see this as justification for our derivation in Section 2.4, which obtained the ISTA block by introducing additional simplifying assumption the distribution of the data at layer L.
Analysis. Following (16), we will consider the natural relaxation of the to "norm" to the l norm, and incorporate a nonnegativity constraint. Consider the objective
< Z|| + X4z20;
log det a+az%z)
(77)
R(2)
where Z E d and a d/NE2, and X(Z20) denotes the characteristic function for the set of elementwise-nonmeaatve matrices Z. As in Appendix A.2, we calculate
V R Z) = aZ (I + aZ*Z)=1.
(78)
1hss is a natural assumption in transformer-type architectures such as CRATE due to the use of LayerNorm blocks although these blocks (indeed, we use them in CRATE) include trainable mean and scale offsets as well an additional mean subtraction operation [64], they initialized to have mean and unit norm, initialization.
hence this assumption corresponds an analysis of the network
This assumption is without loss of generality, as we will see in the analysis below. The reason is that Z Z and Z Z have the same nonzero eigenvalues regardless of the shape of Z. which implies that log det (T + particular, interpreting the norms appropriately (with slight abuse of
az Z)- log det [ + a*)). notation), we have ((Z) Z)) = sp( ZZ", so the purposes of analysis we can always proceed as though Z is tall matrix (as long as we do not use any special properties of a in our derivation).
This assumption is strictly stronger than the previous one, and strictly stronger than an assumption of incoherence on the columns. It corresponds to the representation Z 1+2 being non-collapsed, which we expect hold initialization due the projections U(K] being random.
21

---

We consider an incremental optimization scheme for the highly nonlinear and nonconvex objective p. Following Section 2.3, we optimize locally at a "post-compression"" iterate z We follow the standard proximal majorize-minmmizz framework [70] for incremental/local optimization: this begins with the second-order Taylor expansion for the smooth part of y in a neighborhood of the current iterate Z 1/2
R(Z)= ((((+1/2) + (1 --
(79)
+ S (1 10(2 --
22t1/2 R(z)) (z-z411/a) de,
where for any Z E Rd Z: tZ 1/2 + (1 (1-t) Z) The proximal majorization majorization-miniiizttioo minimization approach alternates two steps to minimize yp:
1. First, use assumptions on z to derive upper bound the operator norm of the Hessian V²R( ( over the effective domain of the optimization problem. We will write for this (uniform) upper bound. This yields a quadratic upper bound for the smooth part of the objective so
2. Then, alternately minimize the smooth part of the quadratic upper bound as a function of Z, and take proximal step the nonsmooth part. It can be shown [70] that corresponds to the iteration
1 V2R(Z) z 2
z = prox & (Il lla (x(20))))
(80)
In the alternating minimization setting of this paper for optimizing (1), we only take one such step, starting at z2+1/2
We will instantiate this program below, showing quantitative error bounds related to assumptions above as necessary. Rather than directly applying the iteration (80), we will derive it below under our aforementioned assumptions.
Starting (79), our first task is to upper bound the quadratic residual. This corresponds estimating
{2 2 z +1/2 V²R(Z,) (z (Z-2z+1/22
- -212)))
(81)
< sup tE|0,1]
V R 2
2 2 + ++//22
(82)
with Cauchy-Schward Using Lemma 5, we can estimate the operator norm term in the previous bound terms of properties of Z 1/2 We need to bound
sup S1
|| (A (I+ (2 (Z, A')))
(I (I
(83)
and Lemma 6 gives that this term is no larger than 9a/4 4 for any Z and any t. With this estimate and (79), we have quadratic upper bound for - 2):
- -R(Z) < - R(Z + V (11/2 2 ze+1/2)
9a +
Z z
241122
(84)
8
F
Mcanwhile, by our assumptions above, we have
-V2 2 (I -a)- 1
112
Z 1+12
(85)
We now minimize the preceding quadratic upper bound as a function of Z. Differentiating, the minimizer Zopt is calculated
1
4 Ze+1/2
Zopt
(86)
9(1 +
and it is well-known that the proximal operator of the sum X(z20) and A|| ||1 is simply the oie-sided soft noft-thresholding operator [70]
prox prox xXI2 .12 (Z) = max Z A1
(87)
22

---

where the maximum is applied elementwise. As in Section 2.4, we may write this elementwise maximum simply as ReLU. Thus, one step of proximal enjoorzation-minimzzzin under our simplifying assumptions takes the form
z = =+ (11+ 4 9(1 + +a)
41 9a:
ze+1/2
(88)
Finally, we point out one additional elaboration which introduces the dictionary D that appears in the
ISTA block in Section 2.4. Notice that for any orthogonal D, one has R(DZ)= R( Z for every Z.
This symmetry implies equivariance properties of V R Z)) and V2 R Z): for every Z and every A and every orthogonal D,
DV =
(89)
DA V R (A)) (A, V2R(DZ) (A)).
(90)
Hence the quadratic Taylor expansion (79) can be written equivalently as
R(Z) = R(D R(D' + ( Z z
(91)
f 1/2 (' (2- (2)) dt.
for any orthogonal D. The significance of this is that we have obtained an expression equivalent to (79), but with z²+1/2 replaced by D' 1'+1/2 moreover, because our approximation arguments above are not affected by left-multiplication of z by an orthogonal matrix (this operation does not change the norms of the columns of z or their correlations, and hence the matrix's incoherence), we can apply exactly the same line of reasoning above to obtain that an equivalent proximal majorization- iitizzatioonmmiimzzattio iteration is given by
4 z e+ = ReLU (1+ 9(1 9(1
41 9a 1
D D' (-112
(92)
for any orthogonal dictionary D. This gives an update quite similar the ISTA block (17) in the case where the dictionary used in Section 2.4 is orthogonal, but without a skip connection.
We thus obtain natural white-box version of this part of the architecture, along with the natural interpretation that its purpose is to sparsify the compressed tokens 222 in (learnable) dictionary. which accords with recent empirical studies [79].
Other architectures? As we mentioned the start of this section, the preceding derivation is performed in the most elementary possible setting in order to demonstrate the majorization- minimization approach for layer design. More precise approximations or assumptions may lead to superior layer designs that better optimize the target objective (1) (and in particular (13)). We mention two here:
1. Beyond -inohee-n features: our derivations above assumed that the incoming representations z were already maximal for the expansion term R in (13). It is desirable to obtain 'perturbative' derivation, which applies in cases where Z 1-2 is not fully orthogonal, but instead oetroorth, in particular incoherent [70]. The derivations above can be adapted to this setting; the perturbation bounds become slightly more delicate, and the ultimate layer (92) changes involve additional normalization.
2. Beyond orthogonal dictionaries: The symmetries of the expansion term R in (13) may be followed lead to pair of dictionaries D and D' and an objective that sparsifies DZD D'. This type of transformation is suggestive of popular architectures that mix over tokens [54, 67], however we consider the simpler form Z in this work. In addition, we have focused for simplicity on orthogonal dictionaries D; as in the previous bullet, one may consider in a similar way dictionaries which are complete and near-orthogonal. Adapting the derivation to overcomplete dictionuries is an interesting future direction that we expect to improve the scalability of CRATE; one avenue to achieve this could be increasing the number of projections UKK] and their embedding dimensions.
23

---

A.3.1 Auxiliary Lemmas
Lemma S. Consider the function
R Z 2
log det (I + +zZzz Z)
(93)
where a > 0 constant. Then we have
Z2R(Z) = aZ (1 + az
(94)
and the Hessian operator R()): RaxN + RdxN satisfies that for any E Rdxa
(22 (A)
(95)
a A (1 a -
(96)
Proof. The gradient calculation follows from [46], for example. For the Hessian, use the usual approach to calculating derivatives: if A is any matrix with the same shape as Z and t >> 0,
θ 0t
V V2 R(Z)
(A)
+
VzR(Z +tA)],
(97)
1=
valid since Ri is smooth. We have
R )
(a + +A) (I a Z (Z+AA))
=( + (A) (I OZ al 2'A A'Z+
=a(Z + +A)
(I I [I aZ 2) 2 + A 4 (I+aZ'2)-
=a + 1A)
(-att)* (I a
2 A' + 2
(I+aZ
where in the fourth line we require that is sufficiently close to 0 in order to invoke the Neumann series. First, notice that the term involving A does not play role in the final expression: after we differentiate with respect
and take limit - O, terms arising due differentiation of t + tA A go to zero, because whenever the summation index & > 0 we have term (-at)* that goes zero as + We thus obtain with the product rule
θ
[t V R [A)]
(98)
Öt
a
AI+ a²
[I
(99)
Lemma 6. One has
9
(I sup ST
+ A'ZI) (I (I+ oZ,
(100)
4
Proof. Fix A satisfying
| < 1. By the triangle inequality,
A aZ Z, Z,)
Z A A Z,) (I+ Z,) -
1|
(101)
A + a - ||p +a Z, [I ZZ - A Z.) + a, -
(102)
For the first term, we note that
A + Z.) 11 ||([I & 1) vec vec(A) (A)
(103)
and since (I a < I, we obtain from Sauchy-Schwarz"
A Zz) S || Ir.
(104)
13 Recall that the eigenvalues of Kronecker product of symmetric matrices are the tensor product of the eigenvalues (with multiplicity).
24

---

We can use similar idea to control the second term. We have from the triangle inequality
( + (Z; + A 2 + 2,
(105)
$ |z A aZiZ,)~
(106)
+1(0+02;22==
22 -2;
(107)
For the first term, have
Z oZ; - + OZ, 2.) 1|
(108)
2.) Z aZ Ze) Zi) vea( A
(109)
$ max (1I + aa Z,) 0 max (Z + aZ Z,) -1 zz | F
(110)
A JF.
(111)
The last estimate follows from computation using the SVD of Z. Mcanwhile, we have for the second term by similar argument (using the fact that the singular values of A and A are identical for any matrix A)
[I az) - S max ((1 + -
zi)
|
(112)
A
(113)
4a
where once again the estimate follows from computation involving the SVD of Z₁ (together with the fact that the function a + a/(1 + ao?) is bounded on 2 0 by (2))))) Putting together, we have obtained
9 S
A - 2, a (Z + o
1 A
(114)
4
which gives the claim after taking suprema.
25

---

B Additional Experiments and Details
In this section, we provide details about our experiments, and report the results of additional experi- ments that were not covered in the main text. CRATE takes arguably the most basic design choices possible, and so we do not attempt to directly compete with state-of-the-ar performance from heavily engineered and empirically designed transformers. The results of our experiments meant to convey a few core messages:
Despite not being engineered to compete with the state-of-the-art, CRATE performs strongly on large-scale real-world datasets, including classification on ImageNet- IK CRATE also achieves strong transfer learning performance.
Because our model is designed through unrolled optinization of a well-understood objective, each layer is interpretable. particular, we can analyze the performance of CRATE, as well as design network modifications, on layer-wise basis. This is powered by arguably unparalleled level of insight into the role of each operator in our network.
We make the simplest possible choices during the design of CRATE, but these can be changed easily while keeping the same framework. We study a few modifications later in this section (Appendix B.4) and show that they do not significantly hurt empirical performance, but emphasize here that there is significant potential for improvement with different architecture choices (and in particular a different theoretical analysis).
B.1 Implementation details
In this subsection, we provide more details for implementing CRATE on vision tasks.
B.1.1 Architecture of CRATE
Architectural modifications. Compared the conceptual architecture proposed in Sections 2.5 and 3, we make the following change for the sake of implementation simplicity:
In the compression step, replace the term Nor [UI,... Uk] in the MSSA operator with another trainable parameter W E RdxpK, Thus the MSSA block becomes
SSA(Z U1)
MSSA U[K], W) = w
(115)
SSA Z Uk)
PyTorch code for CRATE. We provide PyTorch-style code for implementing our proposed network architecture. Algorithm 1 defines the overall architecture, Algorithm 2 and Algorithm 3 contain details for the transformer block, self-attention block (MSSA-block), and MLP block (ISTA-block).
B.1.2 Training Setup
Pre-training on ImageNet-1K. We apply the Lion optimizer [73] for pre-training both CRATE and ViT models. We configure the learning rate as 2.4 x 10 weight decay as 0.5, and batch size as 2,048. We incorporate warm-up strategy with a linear increase over 5 epochs, followed by training the models for a total of 150 epochs with cosine decay. For data augmentation, we only apply the standard techniques, random cropping and random horizontal flipping, on the ImageNet-1K K dataset. We apply label smoothing with smoothing parameter 0.1. One training epoch of CRATE Base takes around 240 seconds using 16 A100 40GB GPUs.
Fine-tuning, We fine-tune our pre-trained CRATE and VIT models on the following target datasets: CIFAR CAAARO/CIFAR// 100 [10], Oxford Flowers- 102 [7]. -xford-III--Pess [16]. We also evaluate our pre-trained models on the commonly used ImageNet Real [36] benchmark. For each fine-tuning task, we use the AdamW optimizer [26]. We configure the learning rate as 5 x 10 weight decay as 0.01, and batch size to be 512. To allow transfer learning. we first resize our input data to 224. For data augmentations, we also adopt several standard techniques: random cropping, random horizontal flipping, and random augmentation (with number transformations n. and magnitude of transformations
14). 14
M "https: h/tttt:://ithub... http::/ggtthuuu.oogggiiiiiieieee augment. py
image-modelsse htslll///lllbmmaa htt/mmain///ee meeiii/iiei timm/ hat//aaaaaaaaa tatta/uto_
26

---

Algorithm 1: PyTorch-style pseudocode CRATENetwork
Class ViT_dictionary definition
CRATE:
# initialization
def init(self, image_size, patch_size, num_classes, dim, depth, heads, mip_dim, pool 'cls', channels 3, dim_head 64, dropout 0., emb_dropout 0.):
# define patch, image dimensions and number of patches image_ height, width pair (image_size)
patch_height, patch_width pair(pacchssize)
num_patiches (image_height patch_height) (image_vidth // aatch_width) patch_dim channels
hatch_ * satch_width
define patch embedding, positional embedding, dropout, transformer self.to_pp talf.to_patt wtt.toooccc_eeeeeeeeeee
Sequential (Rearrange, LayerNorm(patch..ii
Sinear(patch ppatch_dim, dim), LayerNorm(dingg (dim))
self pos_ppp_eaaeep enbeddinn = Parameter (random(1, num_patches 1, dim))
self. cls_token Parameter (random(1, 1, dim))
self dropout Dropout (emb_ (embdrppuut))))
self, transformer Transformer (dim, depth, heads, dim_ mlp_dim, dropout)
define pooling, latent layer, head self self.pool pool self self.to_latent latentt = Identity()
self alffmllphead Sequential (Layer)ormm(imm) Linear (dim, classes))
forward pass def (sel, img)
x = self wellf.t_ppateee
b, n, htt...oo___ccceeeeeeninngggggg (xape
cls_tokes takens repeat clken
d',
x concatenate((cls. tokens, x), dim=1)
x += self, _
:n + 1)]
x = self .rroouu(xx)
self transformer (x)
x mean dim 1) self 'mean' else x(:,
x telf.to_latent.ee
return self het.mmmppeedd
Algorithm 2: Pytorch Style Pseudocode for Transformer Block in CRATE
Class Transformer definition class Transformer:
# initialization
def init(self, dim, depth, heads, dia_head, mlp_dim, dropout = 0.)
# define layers
self, layers []
self. = depth
for in range(depth)
self dropout))])
wpw.nddpLaaerrrooonotttt (LayerNorm(ddin (dim, Attention (dim, heads, dia_head,
self www.lllyyaapppp... append (Layyerloooo (dim, FeedForward mlp_dim, dropout))])
forward pass
def fsrlf, x):
for attn, ff in self selfrlayers: ax) + = fx
return
27

---

Algorithm 3: Pseudocode for Attention and FeedForward
Class FeedForward definition
class FeedForward:
# initialization
def init( dim, hidden_dis dim, dropout - 0., step_size-0.1, lambd-0.11 .1 self.weigee Parameter (Tensor (dim, dim))
knit.kaiming.unik www.iiinnnuiioommmmmmm (self, weight)
self. self.step_size step_size self. lambd
# forward pass
def (self x):
linear self.weight, bias-None)
grad_1 linear (x1, self wei..weighttt0 bias-None)
grad_2 = linear (x, self welgtweight.tui t(, bias-None)
update self self.step_size (grad_2 grad_1) self.step_siee
self.lambd
output relu(x grad_update)
return output
Class Attention definition
class Attention:
initialization
def init(self, dim, heads 8, dim_head 64, dropout 0.): inner_dit dim_head * heads
project_ out not (heads == and dim_head == dim)
self heads heads self, scale head -0.5
self, attend Softmax (dim -1)
self dropout = Dropout (dropout) self qkv = Linear (dim, enner_dirt bias-False) self.to_ou
Sequential (Linear (innerriin___di dim), propooodrropuut
project_out out else nn. Indenitty))) # forward pass
def forward(self, x):
d) b d', self belf.heads) dots matmul v. t.transpose -2)) self.scale
rearrange (self (x)), "b
attn self, self.attendddotta
attn self dropout (attn)
out matmul
rearrange (out, "b
(h a)
return self walf.to_oo to_outttoooooooo (out)
28

---

B.2
Experimental Results
In this subsection, we provide additional experimental results on CRATE, including layer-wise measurements, visualizations, as well as ablation studies.
B.2.1 Layer-wise Evaluation and Visualization
Layer-wise evaluation of compression and sparsity. Similar to Figure 3, we conduct the layer- wise evaluation of compression term and sparsity for CRATE-Tiny, CRATE-Base, and CRATE-Lagge. We observe similar behavior as mentioned in Section 3.1: both the compression term and the sparsity term improves as the layer index increases.
Measure coding layers
Measure outout sparsity arooss layers
Sparsity (ISTA block]
000
602
R-(21) [SSA
802
0.35
0.25
402
20
Layer index f
Layer index !
(a) Compression (Model: CRATE-Tiny).
(b) Sparsity (Model: CRATE- Tiny).
Measure coding rate across layers
Measure output sparsity across layers
1002
Sparsity (ISTA (ISTAbloc) block]
R-12)) [SSA block]
1:00
0.5
2002
802
0.2
602
Layer index
Layer index 1
(c) Compression (Model: CRATE-Base).
(d) Sparsity (Model: CRATE-Base).
Measure coding rate across layers
Measure autput sparsity across layers
2002
Sparsity (ISTA block]
[SSA block]
1003
3002
0.4
1002
0.3
02-
1200
1002
Layer index +f
Layer index
(e) Compression (Model: CRATE-Large).
(f) Sparsity (Model: CRATE-Large).
Figure 5: Left: The compression term (2(+11/2) of the MSSA outputs at different layers. Right: the sparsity of the ISTA output block, || [d/(( N), at different layers.
29

---

Visualizing layer-wise token representations. In Figure 6, we visualize the token representations Z' at different layers & E {I...., 12). We provide more results evaluated on other samples in Appendix B.2.2.
Visualizing layer-wise subspaces in multi-head self-attention. We provide the visualization of UKK [K] in Figure
(=)
(b)
(c) (( 3.
(d) (4
(0) 6.
gg) 7
(h) 8= 8.
(
( 0- 10.
(k) = 11.
() = 12.
Figure 6: Visualizing layer-wise token z representations at cach layer To enhance the visual clarity, we randomly extract 50x 50 sub-matrix from Z' for display purposes. (Model: CRATE-Tiny)
30

---

(a) 1.
2-
(c) == 3.
(d) (- 4.
(e) é = 5.
( =
(g) t= 7.
(h) l = 8,
(i) 8 = 9.
(i) 1= 10.
(k) = 11.
() f 12.
Figure 7: We visualize the [Uf, UK Uf, Uk] E xpp different layers. The (i. (th block in each sub-figure corresponds to (U,) U for i,j E [K] at particular layer f. To enhance the visual clarity, for each subspace U.. we randomly pick 4 directions for display purposes. (Model: CRATE-Tiny)
31

---

B.2.2 Additional Layer-wise Visualization
We provide more results of the layer-wise token representation visualization on other samples in Figure 8, Figure 9, Figure 10, and Figure 11 (Model: CRATE-Base).
Figure 8: Visualizing layer-wise token Z' representations at each layer t. To enhance the visual clarity, we randomly extract 50 sub-matrix from Z' for display purposes. (Sample 1)
Figure 9: Visualizing layer-wise token Z' representations at each layer E. To enhance the visual clarity, we randomly extract 50 sub-mat matrix from z" for display purposes. (Sample 2)
32

---

Lare
Figure 10: Visualizing layer-wise token Z' representations each layer E. To enhance the visual clarity, randomly extract 50 sub-matrix from Z' for display purposes. (Sample 3)
Figure 11: Visualizing layer-wise token Z' representations each layer E. To enhance the visual clarity, we randomly extract 50 sub-matrix from z" for display purposes. (Sample 4)
33

---

B.3 CRATE Ablation
Hyperparameters of CRATE. In Table 2, present evaluation of CRATE trained with various parameters. More specifically, we investigate the effect of number of epochs, weight decay, learning rate, step size (1) and the regularization term (A) in ISTA block. As shown in Table 2, CRATE demonstrates consistently satisfactory performance across a diverse range hyperparameters.
Table 2: Top accuracy of CRATE on various datasets with different architecture design variants when trained on ImageNet.
Model
epoch
weight decay
1r
(ISTA)
À (ISTA)
ImageNet
CRATE-B
150 (default)
0.5 (default)
2.4 x 10"*
0.1
0.1
70.8
CRATE-B
150
0.5
2.4 x 10""
0.02
0.1
70.7
CRATE-B
150
0.5
2.4 10
0.5
0.1
66.7
CRATE-B
150
0.5
2.4 x 10"*
0.1
0.02
70.8
CRATE-B
150
0.5
2.4 10
0.1
0.5
70.5
CRATE-B
90
0.5
2.4 10-1
0.1
0.1
69.5
CRATE-B
10
70.9
CRATE-B
150
1.0
2.4 x 10-4
0.1
0.1
70.3
CRATE-B
150
0.05
2.4 10-1
0.1
0.1
70.2
CRATE-B
150
0.5
4.8 10
0.1
0.1
70.2
CRATE-B
150
0.5
1.2 10
0.1
0.1
70.3
B.4 Exploring Architecture Variants
In this section, we explore the two following alternative architectures. One architecture involves modification to the attention mechanism, while the other involves a modification to the sparsification mechanism. Again, re-emphasize that these choices, although principled, entirely modular and the choices we make here still lead to very simple architectures. A more sophisticated analysis may lead to different, more complicated architectures that perform better in practice. The architectures experiment with are:
caappeession-inspiiee attention mechanism: revert the change in (115). That is, the attention mechanism implements (11) and (12) directly.
Majorization-minniizate proximal step sparsification: instead of (17), implement (92).
We obtain the following classification results in Table 3. After conducting additional simplifications the network architecture (i.e., imposing additional constraints the network architecture design), we discover that CRATE maintains reasonable performance on ImageNet- I.
Table 3: Top accuracy of CRATE on various datasets with different architecture design variants when trained on ImageNet.
Model
MSSA-block
ISTA-block
ImageNet
CRATE-B
default
default
70.8
CRATE-B
Eq. (11) (12)
default
63.3
CRATE-B
default
Eq. (92)
68.6
B.5 Sparse Coding vs. Non-Negative Sparse Coding
In the main body, we used a non-negative sparse coding formulation (16) to obtain the ISTA block as an unrolled proximal gradient step in (17). Suppose that we hadn't done this, and instead directly computed an ISTA block using an unrolled proximal gradient step on (15). Such block would give the following update rule:
zt+1 SN, (Z/11/2
- -DZ
(116)
where Sxn is the soft-threshodingg function
Sx,(x) = sgn(x) ([2] xn)+.
(117)
34

---

applied dleeent-wist to its input matrix. The resulting architecture would be an alternative to CRATE; below, we discuss some empirical and theoretical similarities and differences between the two formulations and architectures.
Empirical evaluation of woft-thresholding-basse architecture. We summarize the results of CRATE with soft-thresholding activation in Table 4. We use À = 10 SXn and set all other hyperparameters the same as in the original CRATE-Base evaluation on ImageNet-1 We find that such soft-thresholding model achieves slightly worse performance-a drop of 3.2% top-1 accuracy comparred to the default CRATE-Basc (with ReLU activation).
Table Top accuracy of CRATE on ImageNet- Ik with different architecture design variants. The soft- thresholding activation Sxn is defined in Eq. (117).
Model
MSSA-block
18TA-block
ImageNet
CIFAR 10*
CIFAR 100*
CRATE-B
default
ReLU activation (default)
70.8%
96.8%
82.7%
CRATE-B
default
soft-thresholding activation
67.6%
96.0%
76.8%
Potential theoretical justification for the performance differential. Previous work [49] studied the phase collapse mechanism for understanding the non-linearities and the convolutional filters used in CNNs such as ResNet [23] on classification tasks. Specifically, they found that replacing the phase collapses with thresholding operators which enforce sparsity largely degrades the classification performance of CNNs. The effect of phase collapse analyzed in [49] is to better separate out the means of different classes within a classification task. This may account in part for the increase in classification accuracy reported in Table 4. On the other hand, we believe that the CRATE architecture will be applicable beyond just classification tasks. In CRATE, the purpose of the training process is to learn the local signal models at each layer (see e.g., Section 2.5). From this perspective, so long as the downstream training task requires semantically meaningful representations of the data distribution, the exact training configuration is of secondary importance. In particular, we may use self-supervised learning methods such as (masked) autoencoding any well-defined notion of class mean. In such cases, a priori we may expect both soft thresholding and nonnegative soft thresholding to perform comparably well. We leave the verification of this to future work.
learn the signal models, whence there may not be
Theoretical justification for non-negative sparse coding. The sparse rate reduction formulation (1) does not include non-negative constraint, and the token representations have marginal distribution equal to mixture of zero-mean Gaussians (which are symmetric around 0). Below, we argue that the non-negative sparse rate reduction optimization and the regular sparse rate reduction optimization engender representations which are qualitatively similar in many ways, which confirms our conceptual understanding of how CRATE performs structured representation learning.
First, we formalize the non-negative sparse rate reduction optimization problem. Let X be the characteristic function (with codomain (0, oc}) of its input proposition. Then the non-negative analogue to (1) is
max AA((ZUUU)) A|| Z x(2 2 0)] FEF
where
2 -(X.
(118)
Although formal analysis of the optimal points of the sparse rate reduction maximization problem (1) or its nonnegative variant (118) is out of scope of this work, we see that the rate reduction maximiza- tion (i.c., / ( Z UK])] has optimal points characterized similarly to [46, Theorem A.6], namely that the representation of each distribution in the mixture is supported on a subspace with nearly isotropic covariance on this subspace, and the supporting subspaces are (nearly) orthogonal. Adding the sparsity term A|| Z lo for some regularizer À would enforce the axis-alignment of the supporting subspaces; when adding in addition the nonnegativity term X Z 2 0). following through the proof of [46, Theorem A. suggests that the argument goes through with suitable modifications (in particular, considering the conclusions for the covariance rather than ZZT), ). This sketch suggests that the statistical and geometric properties of the optimal representation remain the same when adding the non-negative constraint to the sparse rate reduction formulation. We leave detailed proof to future work.
B.6 Pre-training on ImageNet-21K
We inestigate a larger pre-training dataset for training CRATE. Specifically, we first pretrain on ImageNet-21K [9], which contains 14 million images, and then fine-tuned on ImageNet-IK 1.. As
35

---

shown in Table 5, with the CRATE-Baas Base model (22.80M parameters), we achieve 80.2% top- 1 accuracy; this comparable to ViT-Base ( 86M parameters, 83.9%) [40] with around 25% of the parameters. For pre-training on ImageNet-21K, we configure the learning rate to 1 x 10-2, set the weight decay to 0.05, and use a batch size of 4,096. The total number of epochs is 90, with 10 warmup epochs. For fine-tuning on ImageNet-1K, IK, we use the same set of parameters as described in Appendix B.1.2, with the exception of setting the learning rate to 5 10 and having a total of 50 epochs.
Table 5: Top accuracy of CRATE-Base and ViT-Base [40] on various datasets when both models are pre-trained on ImageNet-2Ik. We use the VīT- Base results from [40] basis for comparison.
Model
# parameters
ImageNet
CIFAR 10
CIFAR 100
CRATE-Base
22.80M
80.2%
98.3%
88.3%
ViT-Base [40]
86M
83.9%
99.0%
91.7%
B.7 Evaluating R and sparsity for ViT
We conduct experiments evaluate the R and sparsity of token representations from each layer of a pre-trained ViT-Base (downloaded from hopps://gihhnn con/huggingface! Attorch-image-modell
We summarize the results in Figure 12. We find that without our white- box design, the vanilla ViT does not optimize our proposed sparse rate reduction objective. This contrasts with the results shown in Figures and 4 of the work. wherein we can observe that the compression term R and sparsity value decrease layerwise for CRATE, in accordance with our theory.
1155
Dranne
CMECS
Figure 12: Left: The compression term R (Z(-112) of the multi-head self-attention outputs at different layers. Middle: The sparsity of outputs of the MLP block. || llo/(d (d--dd at different layers. Right: To get more tine-grained understanding of the sparsity of MLP block outputs of ViT. three different thresholds T E (1.0, 0.5, 0.1) and measure E. 1 i-th token representation. (First row model: ViT-Small; second model: ViT-Base; third row model: proposed CRATE-Small).
( < <dd N), where z represents the j-th element in the
36


Reviews:
[{"review_id": "aGDQMjKRzH", "reviewer": "Reviewer_AAKz", "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.", "questions": "No.", "limitations": "No.", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Very interesting claim about transformer.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are particularly grateful for and truly encouraged by your high assessment of our work. Thank you for dedicating your time and expertise to review our paper, and please let us know if you have any additional questions or comments during the discussion period."}, {"review_id": "iEHN6zrhJj", "reviewer": "Reviewer_13ev", "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.", "questions": "1) Please consider larger model with ImageNet-1k top-1 accuracy above 80.0%", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n2) The idea is novel and interesting. Especially, deriving multi-head attention from Maximizing Code Reduction is novel.\n\n3) 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.\n\n2) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and compliments on the idea being “novel and interesting”  as well as the exposition being “well written and easy to follow”.\n\nAs we mentioned in our Public Response, the primary goal of our work is not meant to simply push the state-of-the-art in a particular metric, but rather to demonstrate the promise of the CRATE approach (i.e., white-box deep networks constructed via unrolled optimization). Nevertheless, in addition to achieving this goal, CRATE additionally obtains strong performance and has promising scaling behavior on increasingly larger-scale real world datasets. In particular, per your suggestion (as well as that of Reviewer `YxEu`) to try larger models, we investigated the performance of CRATE on ImageNet-1K when pretrained on ImageNet-21K and fine-tuned on ImageNet-1K. We found that in this setting, CRATE-Base could achieve 80.2% top-1 accuracy, which is comparable to the performance of ViT-Base with around 25% of the parameters -- see our Public Response for more precise details. We will add these experiments into our main tables and explicitly mention that our models perform slightly worse than ViTs in the paper text. Unfortunately, due to time limits, we could not pretrain CRATE-Large on ImageNet-21K; we will add the corresponding empirical results and comparison to ViT-Large in our camera-ready version.\n\nWe hope that the points raised above resolve the doubts you have about this work. Please let us know if you have further questions or comments.\n"}, {"review_id": "XBBw6PVSxa", "reviewer": "Reviewer_vvYz", "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.\nAs 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.\nThe 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.", "questions": "How much does the performance degrade when replacing ReLU by softshrink ? \n\nDo 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 ?\n\nNote that with non-negativity constraints , a subspace becomes a cone. Why would it be important to have a cone rather that a subspace ?", "limitations": "Apart the aspect mentioned below, this is a really strong work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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 \".\nThey 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).\nThe 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). \nThis modelisation work in a tough work, and this paper is a significant contribution.\n", "weaknesses": "To me there is a caveat about the sparse coding hypothesis. \nLine 252:\n\" 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} \"\n\nThis 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] .\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments, and your compliments on the quality of the contribution, the strength of the ideas, and the empirical insights. You bring up a very interesting point that the difference between the sparse coding and non-negative sparse coding formulations may seem important. In the sequel, we will attempt to explain why in fact the choice does not make too much difference conceptually, algorithmically, or empirically. \n\n> *How much does the performance degrade when replacing ReLU by softshrink ?*\n\nRegarding the empirical performance drop when using the regular sparse coding formulation, we report that the CRATE-Base model trains more or less the same (~67.6% top-1 accuracy on ImageNet-1K, which is a drop of 3.2% compared to the ReLU case, using $\\lambda = 10$ and all other hyperparameters the same as in the original CRATE-Base evaluation on ImageNet; c.f., Table 1 in the paper and the rebuttal .pdf file). The results are summarized in Table 2 of the .pdf file. We will add this result to our camera-ready version. The message is that the two networks train comparably well, and one can push performance on either one by more dedicated hyperparameter tuning.\n\n>  *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.*\n\nThank you for bringing the interesting work [1] to our attention. Regarding the specific issue of phase collapse discussed in [1], it is our understanding that the effect of phase collapse which is analyzed in [1] is to better separate out the means of different classes within a classification task. While this may be a cause of the increase in classification accuracy reported above, we believe that our method will be applicable beyond just classification tasks. Indeed, in our framework, we contend that the purpose of the training process is to learn the local signal models at each layer (see e.g., Section 2.5). From this perspective, so long as the downstream training task requires semantically meaningful representations of the data distribution, the exact training configuration is of secondary importance. In particular, we may use self-supervised learning methods to learn the signal models, whence there may not be any well-defined notion of class mean, but such exploration is left to future work. In the camera-ready version, we will expand the discussion to include the work [1] and its discussion of phase collapse along with these clarifications.\n\n> *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-dimensional subspace hypotheses? Gaussians do not care about the sign, isn't it ?*\n\n> *Note that with non-negativity constraints, a subspace becomes a cone. Why would it be important to have a cone rather that a subspace?*\n\nThis is a very good point; if we push a set of input representations through a non-negative sparse coding layer, they will always end up as non-negative, and thus cannot have marginal distributions equal to (an approximation of) a mixture of zero-mean Gaussians, but rather some other distribution. If we propagate this non-negative constraint to the sparse rate reduction problem, we obtain a “nonnegative sparse rate reduction” problem:\n\n$$\\max_{f \\in \\mathcal{F}}\\mathbb{E}[\\Delta R(Z \\mid U_{[K]}) - \\lambda \\\\|Z\\\\|_{0} - \\chi(Z \\geq 0)]$$\n\nwhere $\\chi$ denotes the characteristic function of a set and the algebraic definition of $\\Delta R(Z \\mid U_{[K]})$ (as a linear combination of logdet functions) is given in the paper. Our claim that the CRATE model transforms the data to a mixture of incoherent subspaces stems from the analysis of [2] of the minimizers of the rate reduction objective; in this view, understanding the questions you raise about representations in the presence of nonnegative soft thresholding amounts to whether optimal configurations in our nonnegative sparse rate reduction objective can be understood analogously. We sketch an argument below to this effect.\n\nAlthough formal analysis of the optimal points of the sparse rate reduction maximization problem is out of scope of this work, we see that the rate reduction maximization (i.e., $\\max_{f \\in \\mathcal{F}} \\mathbb{E}[\\Delta R(Z \\mid U_{[K]})]$) has optimal points characterized similarly to [2, Theorem A.6], namely that the representation of each distribution in the mixture is supported on a subspace with nearly isotropic covariance on this subspace, and the supporting subspaces are (nearly) orthogonal. Adding the sparsity term for some regularizer $\\lambda$ would enforce the axis-alignment of the supporting subspaces; when adding in addition the nonnegativity constraint, following through the proof of [2, Theorem A.6] suggests that the argument goes through with suitable modifications (in particular, considering the conclusions for the covariance rather than $Z Z^T$). This sketch suggests that the statistical and geometric properties of the optimal representation remain the same when adding the non-negative constraint to the sparse rate reduction formulation. Since our CRATE model is derived by unrolling this objective, we believe this justifies the conceptual picture we describe in the submission around the ISTA block, although we use ReLU instead of soft thresholding. In the camera-ready version, we will expand the discussion to clarify these conceptual points.\n\nWe again thank you for your detailed and interesting point, and hope we have provided satisfactory responses to your questions. Please let us know if you have further questions or comments.\n\n[1]: Guth, Florentin et al., \"Phase collapse in neural networks.\" arXiv preprint arXiv:2110.05283 (2021).\n\n[2]: Yu, Yaodong et al., \"Learning diverse and discriminative representations via the principle of maximal coding rate reduction.\" Advances in Neural Information Processing Systems 33 (2020).\n"}, {"review_id": "W4WAZJrUDB", "reviewer": "Reviewer_YxEu", "summary": "The authors propose an interpretation of the transformer architecture wherein the component blocks may be interpreted as unrolled optimization steps:\n\n* 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:\n    * $ \\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}^*  )$\n* 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:\n    * $\\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}$\n\nThe relationship between the architecture and the objective is quite involved and I cannot summarise it here.\n\nAs 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.\nSpecifically, they track the sparse coding $R$ mentioned above and the sparsity at each layer in a network.\nDuring training it is then possible to observe the sparse coding rate and sparsity both decrease with the layer index during training.\nIn addition, the network modifications required to match the theory do not appear to reduce performance on ImageNet compared to a similarly sized Vision Transformer.", "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?\n\nWhat do these results imply for transformer design, can you predict anything that is obviously incorrect or valuable just from these results?\n\nCan you design a new layer that performs an additional form of unrolled optimization that is also useful?\n\nIs there a regime of transformer operation that completely fails where this theory would provide a useful insight?\n\nI 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.", "limitations": "The limitation of trying to explain why a given deep learning architecture works is:\n\n1. A method becomes popular because it is surprisingly effecitive\n2. A theory is constructed to explain why it works\n3. The theory demands some small change to the network which only slightly reduces performance\n4. The theory does not immediately permit an improvement to the method that is valuable\n\nI 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.\nIt could be that they are, but I do not see this argument made clearly enough in the paper.\n\nAlternatively, maybe there is some architectural or training improvement that this theory implies that would be extremely valuable.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\nIf this goal is achieved by this paper then it is a significant work.\n\nThe strengths of this paper are therefore mainly found in Figures 3 and 4, which show the statistics $R^c$ and sparsity.\nFigure 3 compares train versus validation while Figure 4 shows how these statistics change during training.\nIt is interesting to observe these values decreasing as the signal travels through the network as this approximately matches what is predicted by the theory.\nThese are the original empirical observations of the paper and they support the paper's claims.\n\nThe motivation and experimental results are both clearly stated.\nSection 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.\nThe 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\n\n* Legibility of the derivation: the derivations in Section 2 are dense and difficult to follow\n    * 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\n    * 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\n    * 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\".\n    * 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\n    * 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?\n* Additional experimental results would be worthwhile:\n    * 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\n    * 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\n    * 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\n    * Comparing these statistics to activation norms would be a good comparison", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. Below, in response to the issues you have raised, we briefly reiterate how we view our work’s core motivations and contributions, then follow with precise responses to specific points raised in your review. Unfortunately due to space constraints we cannot answer all points in full detail, but we are glad to answer any follow-ups in the discussion phase.\n\n## General Comments\n\nYour review suggests that our work’s goals and contributions are wrapped up in providing an interpretation of the popular transformer architecture, and as a byproduct obtaining implications for the design of transformer/self-attention architectures. We agree that this is an excellent path to achieving high-impact research in this area. However, we do not believe it is the only way. We would like to reiterate that the goal of our work *is not to directly explain the existing transformer, but to introduce a new white-box alternative whose operational characteristics are both different and transparent in a useful way.* As we describe in Section 2 of the submission, the model we propose is designed around the principle of learning a representation for nonlinear, multimodal data by incrementally transforming it to a standardized form. Conceptually, we derive our architecture to achieve this goal following an unrolled optimization perspective, leading to a derivation where we understand the role played by each parameter (i.e., a “white-box architecture”). Experimentally, we demonstrate that for a supervised classification task, the characteristics of the learned model agree with its white-box design, and performance is not sacrificed much either, in terms of both accuracy and scaling. Crucially, experiments on pretrained ViTs do not show the same characteristics as our white-box model, suggesting our white-box architecture has novel and highly interpretable characteristics, some of which are visualized in the Appendix (see Figures 6 to 11), that will be useful for network design and analysis beyond classification. \n\n## Responses to Individual Points\n\n### Results training on Imagenet-21k\n\nPlease see the discussion in the general response.\n\n### Computing $R^c$ and sparsity in public pretrained models vs CRATE (the value of Figures 3 and 4)\n\nThank you for your insightful suggestion. We have conducted new experiments to evaluate the $R^c$ and sparsity of token representations from each layer of a pre-trained ViT-Base (downloaded from the `timm` GitHub repo). We summarize the results in Figure 1 of the uploaded rebuttal .pdf file. 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). We will add these results and this comparison to Section 3 of our main body in our camera-ready version. \n\nOverall, this shows that the CRATE model, while looking and performing quite similar to a transformer in experiments, is fully interpretable through the perspective of sparse rate reduction, in a manner that is distinct from the black-box ViT. This is a significant advance in the development of layer-wise interpretable networks -- this is the first, to our knowledge, to achieve performance comparable with standard models such as ViT on ImageNet. \n\n### Implications of our results for transformer design\n\nBoth our theoretical and empirical results suggest that it is possible to design transformer-like architectures from the principle of unrolled optimization. One valuable design choice from our work is that QKV heads in attention may not be necessary, and this could help reduce the number of model parameters (indeed, our design, instantiated in CRATE-Base, has about 25% of the parameters of ViT-Base while demonstrating comparable performance) and make the whole network more efficient. Furthermore, both the MSSA block and the ISTA block are simpler and more interpretable than the existing multi-head attention block and MLP block, and we hope such a minimal and functional network architecture could further improve our understanding of transformer-like architectures.  \n\n### Unrolled optimization, class labels in the derivation\n\nOur CRATE architecture is derived from the goal of transforming the data distribution to a structured form, rather than a specific task. This is why the conceptual derivation in Section 2 does not include any discussion of labels: although labeled samples may be used to learn the parameters of the CRATE model, so long as the downstream training task requires semantically meaningful representations of the data distribution, the exact training configuration is of secondary importance. Intuitively, once we could identify the low-dimensional representations of the high-dimensional data (ie, via optimizing the sparse rate reduction), such representations are effective for classification problems.\n\nWe believe that the unrolled optimization perspective is a major potential advantage of our white-box framework: in settings where, for example, prior information is available about the data distribution (for example, in medical imaging applications or other problems with scientific data), our white-box design allows these structures to be incorporated in a transparent way. \n\n### Comments on the presentation\n\nThank you for your helpful comments. We will incorporate your suggestions in the camera-ready version, using the extra space.\n\nWe hope that the points raised above help clarify the significance of our contributions."}, {"review_id": "3tZy2BJt5D", "reviewer": "Reviewer_tJjQ", "summary": "This paper proposes a optimization target called sparse rate reduction, which is built on previous optimization target called rate reduction [49].  \nBy unrolling the iterative optimization process of sparse rate reduction into neural layers, a transformer-like architecture can be obtained.  \nThe derived white-box transformer-like architecture achieves similar performance with ViT.", "questions": "see above.", "limitations": "see above.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Overall the manuscript is well written and related works are properly cited and discussed, very insightful work.  \nThe 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.  \nThe 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.  \nThe 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.\n\n", "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.  \nAlthough 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.  \n\n1) 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.  \n2) L185 and footnote 6, can you explain more about the content in footnote 6? What is the rigorous math here?  \n3) 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.  \n\nFollowing 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments, and your compliments on the insightfulness of the contribution, the novelty of the idea, the quality of the exposition, and the strength of the empirical results. Below we attempt to resolve the questions you posed.\n\n## ImageNet results versus SOTA\n\nThank you for the comment on the results. Please note that the main goal of this work is to not push the state-of-the-art, and as such the provided results incorporate minimal engineering compared to ViT or any more recent state-of-the-art model. With a more thorough engineering effort, and using the greater understanding of CRATE compared to ViT, one may potentially push the performance of CRATE-like models using the proposed framework beyond the state-of-the-art. Please see the Public Response for details on more experiments that push the performance of CRATE.\n\n## Terseness of presentation\n\nThank you for the suggestion. We will use some of the extra space afforded by the camera-ready version to improve the text, especially incorporating targeted feedback raised by you and other reviewers.\n\n### Footnote 4\n\nThere is some more elaboration on this point in Section 2.5, but we recapitulate it here. The mathematical roles that are separated in this dichotomy are the transformation of the data distribution towards the desired structured form (optimizers of the sparse rate reduction) in the forward pass (which we refer to as “optimization” in the footnote), and the learning of the parameters of these incremental transformations in the backward pass (which we refer to as “learning” in the footnote). Please let us know if this explanation has resolved the confusion; we will make appropriate revisions to the exposition in the paragraph at line 133 in the camera-ready version for improved clarity (merging the footnote into the text).\n\n### Footnote 6\n\nThe “rigorous math” referred to in this footnote refers to the mathematical theory of diffusion models, advanced primarily in [1], captured by two key concepts:\n- Given the score function for the noisy data distribution at a range of noise levels, a diffusion process that ‘follows the score function’ can be used to denoise the input noise distribution towards the data distribution.\n- In fact, it is not necessary to follow a (noisy) diffusion process to generate the data distribution, as there is a fully-deterministic “probability flow” ODE, which also involves the data distribution’s score function, that is mathematically equivalent to the score-following diffusion process. This means that even if one does not add noise when following the score function (as in CRATE), a suitable deterministic process can still transform the input distribution to the data distribution.\n\nWe would be happy to provide more details if necessary in the discussion.\n\n### $\\ell^{0}$ norm explanation\n\nA brief answer is given in Section 2.1, but we reiterate it here. The rate reduction by itself is invariant to arbitrary rotations of the representations. To make the features more efficiently computable, and have human-interpretable structure (e.g., the principal components are just the standard basis vectors), we wish to align the representations with the coordinate axes, so that they become sparse. Thus, we penalize the $\\ell^{0}$ norm, which counts the number of nonzero entries, hence the sparsity, of the representation $Z$. There are several \"relaxations\" of the $\\ell^{0}$ norm which are efficiently optimizable – in Section 2.4 we picked arguably the most basic, namely the $\\ell^{1}$-norm (c.f. LASSO regression, [2]). More choices are possible, and may even yield better performance, but studying this is left to future work.\n\n## Relation between unrolled optimization objective and task; rate reduction and classification\n\nIndeed, the rate reduction framework as introduced in [3] was targeted to the specific case of classification – it uses class labels $\\Pi$ to determine how to group the samples. However, in CRATE, our design of the network architecture (described in Section 2) is completely independent of any labels for the data samples: we derive the architecture from the goal of learning a representation that optimizes the sparse rate reduction objective, and this objective incorporates learnable \"local signal models\" where the labels $\\Pi$ were used in [3]. In our experiments, we eventually learn the parameters of the CRATE model by training on a supervised classification task, but the white-box construction of the model makes the model applicable far beyond the setting of classification. Besides, we also mention that recently rate reduction/information gain type objectives have been used to produce representations which are useful in a variety of tasks, including self-supervised learning [4], generative modeling [5], image segmentation [6], etc. \n\nWe hope that the points raised above resolve the doubts you have about this work, and constitute satisfactory responses to your questions. Please let us know if you have further questions or comments.\n\n[1]: Song, Yang et al., “Score-Based Generative Modeling through Stochastic Differential Equations,” in International Conference on Learning Representations, 2021.\n\n[2]: Wainwright, Martin J. High-dimensional statistics: A non-asymptotic viewpoint. Cambridge university press, 2019.\n\n[3]: Yu, Yaodong, et al., \"Learning diverse and discriminative representations via the principle of maximal coding rate reduction.\" Advances in Neural Information Processing Systems 33 (2020).\n\n[4]: Ding, Tianjiao, et al., \"Unsupervised manifold linearizing and clustering.\" arXiv preprint arXiv:2301.01805 (2023).\n\n[5]: Yang, Allen Y., et al., \"Unsupervised segmentation of natural images via lossy data compression.\" Computer Vision and Image Understanding 110, no. 2 (2008): 212-225.\n\n[6]: Dai, Xili et al., \"Ctrl: Closed-loop transcription to an LDR via minimaxing rate reduction.\" Entropy 24, no. 4 (2022): 456."}]

Author rebuttal:
First, we thank all reviewers for their insightful comments. We are particularly encouraged that reviewers have appreciated:
- The novelty and impact of our central ideas (`13ev`: “...deriving multi-head attention from Maximizing Code Reduction is novel”; `tJjQ`: “provid[ing] a significant extension to [prior work]”; `vvYz`: “...this paper is a significant contribution [to modelisation work]”);
- The benefits of the conceptual framework we have proposed (`vvYz`: “The proposed architecture is conceptually significantly simpler than ViT… and the performance drop is arguably very slight”; `tJjQ`: “The results are promising… the idea of white-box unrolling-based neural network design might be a possible alternative to current black-box design”);
- The quality of the exposition (`YxEu`: “The motivation and experimental results are both clearly stated”; `tJjQ`: “Overall the manuscript is well written and related works are properly cited and discussed, very insightful work”; `13ev`: “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.”);
- The insight presented by our empirical evaluations (`YxEu`: “The results in Table 1… address concerns that this method would reduce performance”; `AAKz`: “Extensive experiments… verify the effectiveness of the proposed method”).

In the remainder of this message, we wish to reiterate our key contributions and address certain concerns raised by the reviewers. In particular, we discuss new empirical results undertaken in response to issues raised by reviewers around CRATE’s scaling behavior, where we demonstrate ImageNet-1K accuracy above 80% after pretraining our CRATE-Base model on ImageNet-21K; these results are presented in the attached .pdf and discussed in full detail below.

## Key Contributions

Our central contribution is that we introduce a new transformer-like architecture (named CRATE), where each network layer/operator is constructed _ab initio_ from the principles of data compression and representation learning. This provides a clear and principled mathematical interpretability to transformer-like networks, by revealing the functions of each network layer while removing unnecessary redundancy from previous empirically designed transformers. In addition, we have shown through experiments that this cleaner and simpler architecture is competitive in performance with the base transformer models (such as ViT) in large-scale real-world vision tasks (e.g. classification on ImageNet). Empirical evaluations further confirm that the overall learned deep networks and their layers clearly perform the mathematical functions they were designed for, i.e., reducing the coding rate and sparsifying the learned representation. We believe this work has shown the promise of eventually bridging the gap between theory and practice of (transformer-like) deep networks.

## Comparison with SOTA; New Experimental Results
Of course, further improving the performance and demonstrating the potential of such new principled models is important -- we agree with several of the reviewers on this point. In this work, however, our goal is not to push the state-of-the-art per se, but rather to develop a more clear and systematic understanding of the extremely ubiquitous transformer-like deep network architectures, by developing a white-box model in this family of architectures. Several reviewers have, nevertheless, suggested some additional experimental improvements so that we may fairly compare to the ViT within more realistic regimes of data and compute; we wish to broadcast the results here. Most notably, as suggested by Reviewers `13ev` and `YxEu`, we have further scaled up the CRATE models. In particular, we pretrained on ImageNet-21K and fine-tuned on ImageNet-1K. As shown in Table 1 of the uploaded pdf file, with the CRATE-Base model (22.80M parameters), we achieve 80.2% top-1 accuracy; this is comparable to ViT-Base (~86M parameters, 83.9%) [1] with around 25% of the parameters. Here, we provide details about the experiments mentioned above. For pre-training on ImageNet-21K, we configure the learning rate to 1e-4, set the weight decay to 0.05, and use a batch size of 4,096. The total number of epochs is 90, with 10 warmup epochs. For fine-tuning on ImageNet-1K, we use the same set of parameters as described in Appendix B.1.2, with the exception of setting the learning rate to 5e-5 and having a total of 50 epochs. Unfortunately, due to time limits, we could not pretrain CRATE-Large on ImageNet-21K; we will add the corresponding empirical results and comparison to ViT-Large in our camera-ready version.

We again thank the reviewers for their insight and hope for a continually enlightening discussion period.

[1]: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby.  ICLR 2021.


Meta-review: Reviewers appreciated the paper's novelty in connecting transformers with Maximizing Rate Reduction, noting its well-structured content and experiments. Strengths included its connection of information-gain frameworks to transformer architecture and verification on real-world datasets. Concerns revolved around the model's performance on ImageNet-1k compared to ViT models and the clarity of dense derivations. The sparse coding hypothesis and its relation to SoftShrink and ReLU were questioned. The authors properly addressed the questions in rebuttal and appendices. Overall, the paper offers promising insights and acs believe it is good for acceptance. It's highly recommended that the clarifications can be added in the camera ready version.
Decision: Accept (poster)