Title: White-Box Transformers via Sparse Rate Reduction
Venue: NeurIPS 2023
Accepted: yes
Keywords: white-box deep neural networks

representation learning

transformer

sparse coding
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

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 .

1 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 to 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 a representation for high-dimensional structured data, such as text [28, 30, 37], images [40, 75], and other types of signals [48, 57]. After the first block, which converts each data point (such as a text corpus or image) into a set or sequence of tokens , further processing is performed on the token sets, in a medium-agnostic manner [28, 40]. A 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

compression

sparsication

Multi-Head Subspace

Self-Aention

(

MSSA

)

Sparse Coding Proximal Step

(

ISTA

)

architecture is empirically designed and lacks a 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 representations 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 in 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-informative pretext task (such as 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-based 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 forward-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, it 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 a more unified framework for designing transformer-like 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 1. 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

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 an 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 multihead 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 interpretable) transformer architecture called CRATE (i.e., Coding RAte reduction TransformEr), 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 on par with much more heavily engineered transformer networks (such as ViT) on a wide variety of tasks (e.g., classification and transfer learning).

2 Technical Approach and Justification

2.1 Objective and Approach

We consider a general learning setup associated with real-world signals. We have some random variable X = [ x 1 , . . . , x N ] ∈ R D × N which is our data source; each x i ∈ R D is interpreted as a token 1 , and the x i 's may have arbitrary correlation structures. We use Z = [ z 1 , . . . , z N ] ∈ R d × N to denote the random variable which defines our representations. Each z i ∈ R d is the representation of the corresponding token x i . We are given B ≥ 1 i.i.d. samples X 1 , . . . , X B ∼ X , whose tokens are x i,b . The representations of our samples are denoted Z 1 , . . . , Z B ∼ Z , and those of our tokens are z i,b . Finally, for a given network, we use Z ℓ to denote the output of the first ℓ layers when given X as input. Correspondingly, the sample outputs are Z ℓ i and the token outputs are z ℓ i,b .

Objective for learning a 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 ∈ R D × N → Z ∈ R d × N which transforms input data X ∈ R D × N with a potentially nonlinear and multi-modal distribution to a (piecewise) linearized and compact feature representation Z ∈ R d × N . While the joint distribution of tokens ( z i ) N i =1 in Z may be sophisticated (and taskspecific), we further contend that it is reasonable and practical to require that the target marginal distribution of individual tokens z i should be highly compressed and structured, amenable for compact coding. Particularly, we require the distribution to be a mixture of low-dimensional (say K ) Gaussian distributions , such that the k th Gaussian has mean 0 ∈ R d , covariance Σ k ⪰ 0 ∈ R d × d , and support spanned by the orthonormal basis U k ∈ R d × p . We denote U [ K ] = ( U k ) K k =1 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 ∆ R ( Z ; U [ K ] ) = R ( Z ) -R c ( Z ; U [ K ] ) , where R and R c are estimates of lossy coding rates to be formally defined in (7) and (8). This also promotes token representations z i from different Gaussians to 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

representation space. 2 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:

where the ℓ 0 norm ∥ Z ∥ 0 promotes the sparsity of the final token representations Z = f ( X ) . 3 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, 70]. Hence it is natural to take an approximation approach that realizes the global transformation f optimizing (1) through a concatenation of multiple, say L , simple incremental and local operations f ℓ that push the representation distribution towards the desired parsimonious model distribution:

where f 0 : R D → R d is the pre-processing mapping that transforms input tokens x i ∈ R D to their token representations z 1 i ∈ R d .

Each incremental forward mapping Z ℓ +1 = 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 a 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 role of each layer as 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 an 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 [ K ] , we opt for a two-step alternating minimization process with a strong conceptual basis: first in Section 2.3, we compress the tokens Z ℓ via a gradient step to minimize the coding rate term min Z R c ( Z ; U [ K ] ) ; second, in Section 2.4, we sparsify the compressed tokens, with a suitably-relaxed proximal gradient step on the difference of the sparsity penalty and the expansion term, i.e., min Z [ λ ∥ Z ∥ 0 -R ( Z )] . Both actions are applied incrementally and repeatedly, as each f ℓ 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 A.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-like operators arise in many contexts. Assume that N = 1 , and the single token x is drawn i.i.d. from an unknown mixture of Gaussians ( N ( 0 , Σ k )) K k =1 supported on low-dimensional subspaces with orthonormal bases U [ K ] = ( U k ) K k =1 and corrupted with additive Gaussian noise w ∼ N ( 0 , I ) , i.e.,

where z is distributed according to the mixture. Our goal is simply to transform the distribution of the noisy token x to the mixture of low-dimensional Gaussians z . Towards incremental construction of a representation f for this model following (2), we reason inductively: if z ℓ is a noisy token (3) at noise level σ ℓ , it is natural to produce z ℓ +1 by denoising at the level σ ℓ . In the mean-square sense, the optimal estimate is E [ z | z ℓ ] , which has a variational characterization (e.g. [12]):

Setting z ℓ +1 = E [ z | 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 x ↦→ q ℓ ( x ) denote the density of z ℓ , Tweedie's formula [13] allows us to express the optimal representation solving (4) in closed-form:

Tweedie's formula expresses the optimal representation in terms of an additive correction (in general a nonlinear function of z ℓ ) to the noisy observations by the gradient of the log-likelihood of the distribution of the noisy observations, giving the optimal representation a 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 ∇ x log q ℓ , which, when combined with (5) and some technical assumptions 4 , gives the following approximation (shown in Appendix A.1). Let ⊗ denote the Kronecker product; then we have

This operation resembles a self-attention layer in a standard transformer architecture with K heads, sequence length N = 1 , the 'query-key-value' constructs being replaced by a single linear projection U ∗ k z ℓ 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-type layers in the transformation f . Given an initial sample x following the model (3), we can repeatedly apply local transformations to the distribution with (6) in order to realize the incremental mapping f : x → 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 in a transformer resembles the scorematching operator that aims to transform a token z ℓ 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-for example, co-occurrences between words or object parts in language and image data (resp.)-which we should also learn. Thus, we should compress / denoise / transform such a set of tokens together. To this end, we need a measure of quality, i.e., compactness, for the resulting representation of the set of tokens. A 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 a closed form. If we view the tokens in Z ∈ R d × N 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:

In practice, the data distribution is 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 U [ K ] = ( U k ) K k =1 , where U k ∈ R d × p . Although the joint distribution of the tokens Z is unknown, the desired marginal distribution of each token z i is a

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:

We would like to compress (or denoise) the set of tokens against these subspaces by minimizing the coding rate. The gradient of R c ( Z ; U [ K ] ) is

The above expression approximates the residual of each projected token U ∗ k z i regressed by other tokens U ∗ k z j [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 ∗ k Z ) ∗ ( U ∗ k Z ) and convert it to a distribution of membership with a softmax, namely softmax(( U ∗ k Z ) ∗ ( U ∗ k Z )) . Then, as we show in Appendix A.2, if we only use similar tokens to regress and denoise each other, then a gradient step on the coding rate with learning rate κ can be naturally approximated as follows:

where MSSA is defined through an SSA operator as:

Here the SSA operator in (11) resembles the attention operator in a typical transformer [28], except that here the linear operators of value, key, and query are all set to be the same as the subspace basis, i.e., V = K = Q = U ∗ k . 6 Hence, we name SSA ( · | U k ) : R d × N → R p × N the S ubspace S elfA ttention (SSA) operator (more details and justification can be found in (72) in Appendix A.2). Then, the whole MSSA operator in (12), formally defined as MSSA ( · | U [ K ] ): R d × N → R d × N and called the M ulti-Head S ubspace S elfA ttention (MSSA) operator, aggregates the attention head outputs by averaging using model-dependent weights, similar in concept to the popular multi-head self-attention operator in existing transformer networks. The overall gradient step (10) resembles the multi-head self-attention implemented with a skip connection in transformers.

Notice that if we have N = 1 tokens as well as take an aggressive gradient step ( κ = 1 ) and tune the quantization error ( ϵ = √ p/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 a linear function here, while in (6) it is done by a nonlinear mixture-of-experts type function. 7 This provides two very related interpretations of the multi-head self-attention operator, as denoising and compression against a mixture of low-dimensional subspaces.

2.4 MLP via Iterative Shrinkage-Thresholding Algorithms (ISTA) for Sparse Coding

In the previous subsection, we focused on how to compress a set of tokens against a 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 a more compact and structured (i.e., parsimonious ) representation. From (1) and (7), this term is

Multi-Head Subspace

Self-Aention

(

MSSA

)

Add & LayerNorm

Proximal Step

(

ISTA

)

LayerNorm

MSSA Block

SSA

(head 1

SSA

(head K)

Aggregate

. . .

SSA Block

Autocorrelation

& Softmax

ISTA Block

Activation

where R ( Z ) denotes the coding rate of the whole token set, as defined in (7). In addition to sparsification via the ∥ Z ∥ 0 term, the expansion term R ( Z ) in (13) promotes diversity and noncollapse of the representation, a highly desirable property. However, prior work has struggled to realize this benefit on large-scale datasets due to poor scalability of the gradient ∇ 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 a (complete) incoherent or orthogonal dictionary D ∈ R d × d , and ask to sparsify the intermediate iterates Z ℓ +1 / 2 with respect to D . That is, Z ℓ +1 / 2 = DZ ℓ +1 where Z ℓ +1 is more sparse. The dictionary D is global, i.e., is used to sparsify all tokens simultaneously. By the incoherence assumption, we have D ∗ D ≈ I d ; thus from (7) we have R ( Z ℓ +1 ) ≈ R ( DZ ℓ +1 ) = R ( Z ℓ +1 / 2 ) . Thus we approximately solve (13) with the following program:

The above sparse representation program is usually solved by relaxing it to an unconstrained convex program, known as LASSO:

In our implementation, motivated by Sun et al. [33] and Zarka et al. [35], we also add a non-negative constraint to Z ℓ +1 ,

which we then incrementally optimize by performing an unrolled proximal gradient descent step, known as an ISTA step [8], to give the update:

In Appendix A.3, we will show one can arrive at a similar operator to 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:

(Sections 2.2 and 2.3) Local denoising and compression of tokens within a sample towards a mixture-of-subspace structure, leading to the multi-head subspace self-attention block MSSA ;

(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-based transformer layer, illustrated in Figure 2,

Composing multiple such layers following the incremental construction of our representation in (2), we obtain a white-box transformer architecture that transforms the data tokens towards a compact and sparse union of incoherent subspaces.

This model has the parameters ( U ℓ [ K ] ) L ℓ =1 and ( D ℓ ) L ℓ =1 , which are learned from data via backpropagation . Notably, in each layer ℓ , the learned U ℓ [ K ] retain their interpretation as incoherent bases for supporting subspaces for the mixture-of-Gaussians model at layer ℓ , and the learned D ℓ retains its interpretation as a sparsifying dictionary at layer ℓ . We emphasize that the parameters U ℓ [ K ] and D ℓ are dependent on the layer ℓ -that is, we learn a different set of parameters at each layer. This is because at each layer we 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 at 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 at each layer, denoise/compress/sparsify the input) and the backward pass (learn the local signal models from data via supervision).

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 TransformEr), 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 transformer 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 is 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 if layers of the learned network indeed perform their design objective-say performing incremental optimization for the objective (1). Second, despite their simplicity, our experiments will actually reveal the vast practical potential of our so-derived 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.1. 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-Tiny, CRATE-Small, CRATE-Base, and CRATE-Large. A PyTorchstyle pseudocode can be found in Appendix B.1, which contains more implementation details. For training using supervised classification, we first take the CLS token z b = z L +1 1 ,b of for each sample, then apply a linear layer; the output of this linear layer u b . = Wz b is used as input to the standard cross-entropy loss. The overall loss averages over all samples b ∈ [ B ] .

Datasets and optimization. Wemainly consider ImageNet-1K [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 ImageNet-1K as pre-trained models, we fine-tune CRATE on several commonly used downstream datasets (CIFAR10/100, Oxford Flowers, Oxford-IIT-Pets). More details about the training and datasets can be found in Appendix B.1.

2

4

6

8

10

12

Layer index -

600

700

800

900

1000

1100

R c

(

Z

) [SSA block]

Measure coding rate across layers

train

val

2

4

6

8

10

12

Layer index -

0.20

0.25

0.30

0.35

0.40

0.45

0.50

Sparsity [ISTA block]

Measure output sparsity across layers

train

val

2

4

6

8

10

12

800

900

1000

1100

1200

1300

1400

1500

1600

R c

(

Z

) [SSA block]

Measure coding rate across layers

rand init

epoch 1

epoch 20

epoch 150

2

4

6

8

10

12

0.15

0.20

0.25

0.30

0.35

0.40

0.45

0.50

Sparsity [ISTA block]

Measure output sparsity across layers

rand init

epoch 1

epoch 20

epoch 150

Layer index -

Layer index -

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 to optimize the compression term R c ( Z ) and the ISTA block to sparsify the token representations (corresponding to the sparsification term ∥ Z ∥ 0 ). To understand whether CRATE indeed optimizes these terms, for each layer ℓ , we measure (i) the compression term R c ( Z ℓ +1 / 2 ) on the MSSA block outputs Z ℓ +1 / 2 ; and (ii) sparsity ∥ Z ℓ +1 ∥ 0 on the ISTA block outputs Z ℓ +1 . Specifically, we evaluate these two terms by using training/validation samples from ImageNet-1K. Both terms are evaluated at the per-sample level and averaged over B = 10 3 samples.

Figure 3 shows the plots of these two key measures at all layers for the learned CRATE-small 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 as well as 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-both 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 see the effect of learning, we present the evaluations on 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 training-learning better subspaces U ℓ [ K ] and dictionaries D ℓ -the designed blocks start to optimize the objectives much more effectively.

Visualizing layer-wise token representations. To gain a better understanding of the token representations of CRATE, we visualize the output of each ISTA block at layer ℓ in Figure 6 of Appendix B.2. Specifically, we visualize the Z ℓ +1 via heatmap plots. We observe that the output Z ℓ +1 becomes more sparse as the layer increases. Moreover, besides the sparsity, we also find that Z ℓ +1 becomes

more structured (i.e., low-rank), which indicates that the set of token representations become closer to linear subspaces, confirming our 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 ] matrices used in the MSSA block. In Section 2.3, we assumed that U ℓ [ K ] were incoherent to capture different 'views' of the set of tokens. In Fig. 7 of Appendix B.2, we first normalize the columns in each U ℓ k , then we visualize the [ U ℓ 1 , . . . , U ℓ K ] ∗ [ U ℓ 1 , . . . , U ℓ K ] ∈ R pK × pK . The ( i, j ) -th block in each sub-figure corresponds to ( U ℓ i ) ∗ U ℓ j for i, j ∈ [ K ] at a particular layer ℓ . We find that the learned U ℓ [ K ] are approximately incoherent, which aligns well with our assumptions. One interesting observation is that the U ℓ [ K ] becomes more incoherent when the layer index ℓ is larger, which suggests that the token representations are more separable. This mirrors the situation in other popular deep networks [58].

3.2 Evalutions of CRATE on Large Real-World Datasets and Tasks

We now study the empirical performance of the proposed networks by measuring their top-1 accuracy on ImageNet-1K as well as transfer learning performance on 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-1K 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 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 transformerlike network architectures as incremental optimization schemes to learn compressed and sparse representation of the input data (or token sets). The so derived and learned deep architectures are not only fully mathematically interpretable, but also consistent on a layer-by-layer level with their design 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 as 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.

References

Charles M Stein. 'Estimation of the Mean of a Multivariate Normal Distribution'. The Annals of Statistics 9.6 (Nov. 1981), pp. 1135-1151. 5.

Bruno A Olshausen and David J Field. 'Sparse coding with an overcomplete basis set: A strategy employed by V1?' Vision research 37.23 (1997), pp. 3311-3325. 2.

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.

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.

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. 413429. 2.

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.

Maria-Elena Nilsback and Andrew Zisserman. 'Automated flower classification over a large number of classes'. 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing . IEEE. 2008, pp. 722-729. 26.

Amir Beck and Marc Teboulle. 'A fast iterative shrinkage-thresholding algorithm for linear inverse problems'. SIAM journal on imaging sciences 2.1 (2009), pp. 183-202. 7.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 'Imagenet: A largescale hierarchical image database'. 2009 IEEE conference on computer vision and pattern recognition . Ieee. 2009, pp. 248-255. 8, 35.

Alex Krizhevsky, Geoffrey Hinton, et al. 'Learning multiple layers of features from tiny images' (2009). 26.

Karol Gregor and Yann LeCun. 'Learning fast approximations of sparse coding'. Proceedings of the 27th International Conference on International Conference on Machine Learning . Omnipress. 2010, pp. 399-406. 2.

László Györfi, Michael Kohler, Adam Krzyzak, and Harro Walk. A Distribution-Free Theory of Nonparametric Regression . Springer New York, Dec. 2010. 4.

Bradley Efron. 'Tweedie's Formula and Selection Bias'. Journal of the American Statistical Association 106.496 (2011), pp. 1602-1614. 2, 5, 16.

Martin Raphan and Eero P Simoncelli. 'Least squares estimation without priors or supervision'. Neural computation 23.2 (Feb. 2011), pp. 374-420. 5.

Pascal Vincent. 'A connection between score matching and denoising autoencoders'. Neural computation 23.7 (July 2011), pp. 1661-1674. 5.

Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. 'Cats and dogs'. 2012 IEEE conference on computer vision and pattern recognition . IEEE. 2012, pp. 3498-3505. 26.

Daniel A Spielman, Huan Wang, and John Wright. 'Exact Recovery of Sparsely-Used Dictionaries' (June 2012). arXiv: 1206.5882 [cs.LG] . 2.

Joan Bruna and Stéphane Mallat. 'Invariant scattering convolution networks'. IEEE transactions on pattern analysis and machine intelligence 35.8 (Aug. 2013), pp. 1872-1886. 2.

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.

Singanallur V Venkatakrishnan, Charles A Bouman, and Brendt Wohlberg. 'Plug-and-Play priors for model based reconstruction'. 2013 IEEE Global Conference on Signal and Information Processing . Dec. 2013, pp. 945-948. 5.

Rémi Gribonval, Rodolphe Jenatton, and Francis Bach. 'Sparse and spurious: dictionary learning with noise and outliers' (July 2014). arXiv: 1407.5155 [cs.LG] . 2.

Jascha Sohl-Dickstein, Eric A Weiss, Niru Maheswaranathan, and Surya Ganguli. 'Deep Unsupervised Learning using Nonequilibrium Thermodynamics' (Mar. 2015). arXiv: 1503. 03585 [cs.LG] . 2, 5.

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

René Vidal, Yi Ma, and Shankar Sastry. Generalized Principal Component Analysis . Springer Verlag, 2016. 5.

Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. 'Mask R-CNN' (Mar. 2017). arXiv: 1703.06870 [cs.CV] . 1.

Ilya Loshchilov and Frank Hutter. 'Decoupled weight decay regularization'. arXiv preprint arXiv:1711.05101 (2017). 26.

Yaniv Romano, Michael Elad, and Peyman Milanfar. 'The Little Engine That Could: Regularization by Denoising (RED)'. SIAM journal on imaging sciences 10.4 (Jan. 2017), pp. 18041844. 5.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 'Attention is all you need'. Advances in neural information processing systems 30 (2017). 1-3, 6.

Yubei Chen, Dylan Paiton, and Bruno Olshausen. 'The sparse manifold transform'. Advances in neural information processing systems 31 (2018). 2.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 'Bert: Pre-training of deep bidirectional transformers for language understanding'. arXiv preprint arXiv:1810.04805 (2018). 1.

Tero Karras, Samuli Laine, and Timo Aila. 'A Style-Based Generator Architecture for Generative Adversarial Networks' (Dec. 2018). arXiv: 1812.04948 [cs.NE] . 1.

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.

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

Yang Song and Stefano Ermon. 'Generative Modeling by Estimating Gradients of the Data Distribution' (July 2019). arXiv: 1907.05600 [cs.LG] . 2.

John Zarka, Louis Thiry, Tomás Angles, and Stéphane Mallat. 'Deep network classification by scattering and homotopy dictionary learning'. arXiv preprint arXiv:1910.03561 (2019). 7.

Lucas Beyer, Olivier J Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. 'Are we done with imagenet?' arXiv preprint arXiv:2006.07159 (2020). 26.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 'Language models are few-shot learners'. Advances in neural information processing systems 33 (2020), pp. 1877-1901. 1.

Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. 'End-to-End Object Detection with Transformers' (May 2020). arXiv: 2005.12872 [cs.CV] . 1.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. 'A Simple Framework for Contrastive Learning of Visual Representations'. Proceedings of the 37th International Conference on Machine Learning . Ed. by Hal Daumé Iii and Aarti Singh. Vol. 119. Proceedings of Machine Learning Research. PMLR, 2020, pp. 1597-1607. 1.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. 'An image is worth 16x16 words: Transformers for image recognition at scale'. arXiv preprint arXiv:2010.11929 (2020). 1, 3, 10, 36.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. 'Denoising diffusion probabilistic models'. Advances in Neural Information Processing Systems 33 (2020), pp. 6840-6851. 2.

Zahra Kadkhodaie and Eero P Simoncelli. 'Solving Linear Inverse Problems Using the Prior Implicit in a Denoiser' (July 2020). arXiv: 2007.13640 [cs.CV] . 5.

Jiaming Song, Chenlin Meng, and Stefano Ermon. 'Denoising Diffusion Implicit Models' (Oct. 2020). arXiv: 2010.02502 [cs.LG] . 2, 5.

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). arXiv: 2011.13456 [cs.LG] . 2, 5.

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 processing systems 33 (2020), pp. 6827-6839. 2.

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. 9422-9434. 2, 3, 5, 19, 24, 35.

Yuexiang Zhai, Zitong Yang, Zhenyu Liao, John Wright, and Yi Ma. 'Complete dictionary learning via l 4-norm maximization over the orthogonal group'. The Journal of Machine Learning Research 21.1 (2020), pp. 6622-6689. 2.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. 'Vivit: A video vision transformer'. Proceedings of the IEEE/CVF international conference on computer vision . 2021, pp. 6836-6846. 1.

Florentin Guth, John Zarka, and Stephane Mallat. 'Phase collapse in neural networks'. arXiv preprint arXiv:2110.05283 (2021). 35.

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 [cs.CV] . 1.

Geoffrey Hinton. How to represent part-whole hierarchies in a neural network . 2021. arXiv: 2102.12627 [cs.CV] . 6.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. 'Learning Transferable Visual Models From Natural Language Supervision'. Proceedings 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. 87488763. 1.

Bahareh Tolooshams and Demba Ba. 'Stable and Interpretable Unrolled Dictionary Learning'. arXiv preprint arXiv:2106.00058 (2021). 2.

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.

Kwan Ho Ryan Chan, Yaodong Yu, Chong You, Haozhi Qi, John Wright, and Yi Ma. 'ReduNet: A White-box Deep Network from the Principle of Maximizing Rate Reduction'. Journal of Machine Learning Research 23.114 (2022), pp. 1-103. 2-8, 19, 20.

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.01916 (2022). 2.

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.

Hangfeng He and Weijie J Su. 'A law of data separation in deep learning'. arXiv preprint arXiv:2210.17020 (2022). 10.

Geoffrey Hinton. The Forward-Forward Algorithm: Some Preliminary Investigations . 2022. arXiv: 2212.13345 [cs.LG] . 2.

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.

Frederic Koehler, Alexander Heckett, and Andrej Risteski. 'Statistical Efficiency of Score Matching: The View from Isoperimetry' (Oct. 2022). arXiv: 2210.00726 [cs.LG] . 2.

Yi Ma, Doris Tsao, and Heung-Yeung Shum. 'On the principles of parsimony and selfconsistency for the emergence of intelligence'. Frontiers of Information Technology & Electronic Engineering 23.9 (2022), pp. 1298-1323. 1, 3.

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.09120 (2022). 2.

Mary Phuong and Marcus Hutter. 'Formal algorithms for transformers'. arXiv preprint arXiv:2207.09238 (2022). 21.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 'High-resolution image synthesis with latent diffusion models'. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2022, pp. 10684-10695. 1, 2.

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 J Fleet, and Mohammad Norouzi. 'Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding' (May 2022). arXiv: 2205.11487 [cs.CV] . 1.

Asher Trockman, Devin Willmott, and J Zico Kolter. 'Understanding the Covariance Structure of Convolutional Filters' (Oct. 2022). arXiv: 2210.03651 [cs.CV] . 23.

Rene Vidal. Attention: Self-Expression Is All You Need . Unpublished; available: https: //openreview.net/forum?id=MmujBClawFo . 2022. 2.

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.

John Wright and Yi Ma. High-Dimensional Data Analysis with Low-Dimensional Models: Principles, Computation, and Applications . Cambridge University Press, 2022. 4, 22, 23.

Yongyi Yang, Zengfeng Huang, and David P Wipf. 'Transformers from an Optimization Perspective'. Advances in Neural Information Processing Systems . Ed. by S Koyejo, S Mohamed, A Agarwal, D Belgrave, K Cho, and A Oh. Vol. 35. Curran Associates, Inc., 2022, pp. 36958-36971. 2.

Sitan Chen, Giannis Daras, and Alexandros G Dimakis. 'Restoration-Degradation Beyond Linear Diffusions: A Non-Asymptotic Analysis For DDIM-Type Samplers' (Mar. 2023). arXiv: 2303.03384 [cs.LG] . 5.

Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, et al. 'Symbolic discovery of optimization algorithms'. arXiv preprint arXiv:2302.06675 (2023). 8, 26.

Brent De Weerdt, Yonina C Eldar, and Nikos Deligiannis. 'Designing Transformer Networks for Sparse Recovery of Sequential Data Using Deep Unfolding'. ICASSP 2023 - 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) . June 2023, pp. 1-5. 2.

Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. 'Scaling vision transformers to 22 billion parameters'. arXiv preprint arXiv:2302.05442 (2023). 1.

Benjamin Hoover, Yuchen Liang, Bao Pham, Rameswar Panda, Hendrik Strobelt, Duen Horng Chau, Mohammed J Zaki, and Dmitry Krotov. 'Energy Transformer' (Feb. 2023). arXiv: 2302.07253 [cs.LG] . 2.

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.

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:2302.06015 (2023). 2.

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.

Ravid Shwartz-Ziv and Yann LeCun. 'To Compress or Not to Compress-Self-Supervised Learning and Information Theory: A Review'. arXiv preprint arXiv:2304.09355 (2023). 2.

Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. 'Consistency models'. arXiv preprint arXiv:2303.01469 (2023). 2.

Thus we have

This gives

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 at a slightly more technical level. Connections between denoising and score matching are well-understood [60], and computing the optimal denoising function (i.e., the conditional expectation) against a mixture-ofGaussians model is a rather simple computation giving existing tools such as Tweedie's formula [13]. These are not our main contributions. Instead, the main contributions of Section 2.2 are two-fold:

First, we demonstrate a mechanism to learn representations via denoising within a idealized mixture of Gaussian data model for a single token (i.e., with sequence length N = 1 ).

Second, we illustrate the similarities between a such-derived representation learning scheme and existing self-attention layers within the transformer (with sequence length 1 ), thus demonstrating an interpretation of the self-attention layer as a generalized mechanism to denoise against a mixture-of-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: x ↦→ z ℓ , z ↦→ z ℓ +1 , and σ ↦→ σ ℓ .

Proposition 1. Let u 1 , . . . , u K ∈ R d be independent and have distribution u k ∼ N ( 0 , Σ k ) for Σ k ⪰ 0 , and let z take value u k with probability π k > 0 . Let w ∼ N ( 0 , I d ) be independent of z . Let x . = z + σ w . Let x ↦→ q ( x ) be the density of x . We define

and assume that π i det( M i ) = π j det( M j ) for all 1 ≤ i ≤ j ≤ K . Then we have

where ⊗ denotes the Kronecker product, i.e., the block matrix defined by

Proof. Let u be the multinomial random variable such that z = z u , so that u has probability mass function π . Then by the law of total probability, we have

where q ( x | k ) is the conditional density of x given the event { u = k } . To compute this quantity, note that conditional on the value of u , we have

Putting this all together, we get

Now define M k . = ( Σ k + σ 2 I d ) -1 / 2 . With this notation, we have

Given our assumption that each π k det( M k ) is the same, we have

Now we provide a final justification for the result cited in Section 2.2.

Approximation 2. In the setting of Proposition 1, diagonalize Σ k = U k Λ k U ∗ k where U k ∈ R d × p is orthogonal and Λ k ≻ 0 ∈ R p × p is diagonal. 9 Then we have the approximation

Proof. We have

Now define P k . = I d -σ M k , and let U ⊥ k ∈ R d × ( d -p ) be an orthogonal complement of U k . Then we have

Thus P k is approximately a projection when σ is small. Under this algebraic relation, we have ∇ x log q ( x ) (51)

Plugging this into Tweedie's formula, we have

Remark 3 . Although Approximation 2 is stated as an approximation rather than as a proposition, we believe it should be possible without too much extra work to convert it into a statement of asymptotic equivalence as σ → 0 (in particular, holding for σ below the smallest (nonzero) eigenvalue of any Σ k . Most approximations taken in the derivation of Approximation 2 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 σ → 0 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 to 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 a generalized coding rate function R c ( · ; U [ K ] ) which measures the coding rate with respect to a set of subspaces U [ K ] as opposed to a 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 U ∗ k now, which we interpret as projection onto a 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, we make perhaps the simplest possible choices . As a result, we do not claim that our network is optimally designed, but rather that the principles we develop in this work (compression, denoising, sparsification, unrolled optimization) can provide the backbone for far superior and more interpretable network architectures in the future on sundry tasks. As it is, with our straightforward, simple, and interpretable design, we still obtain meaningful conceptual results and very solid empirical performance.

We now give the derivation of the approximation alluded to in Section 2.3.

Approximation 4. Let Z ∈ R d × N have unit-norm columns, and U [ K ] = ( U 1 , . . . , U K ) such that each U k ∈ R d × p is an orthogonal matrix, the ( U k ) K k =1 are incoherent, and the columns of Z approximately lie on ⋃ K k =1 Span( U k ) . Let γ = p Nϵ 2 . Let κ > 0 . Then

where as in Section 2.3 we have

where softmax( · ) is the softmax operator (applied to each column of an input matrix), i.e.,

Proof. According to (9), the gradient ∇ Z R c ( Z ; U [ K ] ) is

Notice that according to [55], the gradient is precisely the residual of a ridge regression for each (projected) token U ∗ k z i using other projected tokens U ∗ k z j as the regressors, hence being the residual of an auto-regression.

However, as we have seen in the work of ReduNet [55], computing the inverse ( I + γ ( U ∗ k Z ) ∗ ( U ∗ k Z )) -1 can be expensive. Hence for computational efficiency, we may approximate it with the first order term of its von Neumann expansion:

Notice that the term ( U ∗ k Z ) ∗ ( U ∗ k Z ) 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-membership indicator with a softmax operation, whence (69) becomes

Then, we can rewrite the above approximation to the gradient of R c as:

Thus the gradient descent step with learning rate κ > 0 gives

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, we show that the second alternating step - gradient descent on the overall coding rate plus a sparse regularization term - has heuristic connections to a 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 to pose the minimization step as a LASSO was a simple, reliable, and interpretable choice which works well in practice, but is nonetheless not backed up by rigorous theoretical justification. In the following subsection, we provide a mathematical justification for a reformulation of the minimization step using a majorizationminimization framework. We further show that the associated unrolled optimization step bears a strong resemblance to the ISTA step. This confirms our earlier discussion - we took the simplest 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.

Assumptions. In this section, we present a rigorous optimization analysis of an incremental minimization approach to the objective (13). We will show that under two simplifying assumptions, namely

The columns of Z ℓ +1 / 2 are normalized, in the sense that diag(( Z ℓ +1 / 2 ) ∗ Z ℓ +1 / 2 ) = 1 ; 10

We have d ≥ N , 11 and the columns of Z ℓ +1 / 2 are orthogonal, so that ( Z ℓ +1 / 2 ) ∗ Z ℓ +1 / 2 = I . 12

the approach leads to an update iteration that is equal to a slightly simplified version of the ISTA block (17). We see this as a justification for our derivation in Section 2.4, which obtained the ISTA block by introducing an additional simplifying assumption on the distribution of the data at layer ℓ .

Analysis. Following (16), we will consider the natural relaxation of the ℓ 0 'norm' to the ℓ 1 norm, and incorporate a nonnegativity constraint. Consider the objective

where Z ∈ R d × N and α = d/Nε 2 , and χ { Z ≥ 0 } denotes the characteristic function for the set of elementwise-nonnegative matrices Z . As in Appendix A.2, we calculate

We consider an incremental optimization scheme for the highly nonlinear and nonconvex objective φ . Following Section 2.3, we optimize locally at a 'post-compression' iterate Z ℓ +1 / 2 . We follow the standard proximal majorize-minimize framework [70] for incremental/local optimization: this begins with the second-order Taylor expansion for the smooth part of φ in a neighborhood of the current iterate Z ℓ +1 / 2 :

where for any Z ∈ R d × N , Z t = t Z ℓ +1 / 2 +(1 -t ) Z . The proximal majorization-minimization approach alternates two steps to minimize φ :

First, use assumptions on Z ℓ +1 / 2 to derive an upper bound on the operator norm of the Hessian ∇ 2 R ( Z ) over the effective domain of the optimization problem. We will write L for this (uniform) upper bound. This yields a quadratic upper bound for the smooth part of the objective φ .

Then, alternately minimize the smooth part of the quadratic upper bound as a function of Z , and take a proximal step on the nonsmooth part. It can be shown [70] that corresponds to the iteration

In the alternating minimization setting of this paper for optimizing (1), we only take one such step, starting at Z ℓ +1 / 2 .

We will instantiate this program below, showing quantitative error bounds related to our assumptions above as necessary. Rather than directly applying the iteration (80), we will derive it below under our aforementioned assumptions.

Starting at (79), our first task is to upper bound the quadratic residual. This corresponds to estimating

with Cauchy-Schwarz. Using Lemma 5, we can estimate the operator norm term in the previous bound in terms of properties of Z ℓ +1 / 2 . We need to bound

and Lemma 6 gives that this term is no larger than 9 α/ 4 for any Z and any t . With this estimate and (79), we have a quadratic upper bound for -R ( Z ) :

Meanwhile, by our assumptions above, we have

We now minimize the preceding quadratic upper bound as a function of Z . Differentiating, the minimizer Z opt is calculated as

and it is well-known that the proximal operator of the sum of χ { Z ≥ 0 } and λ ∥ · ∥ 1 is simply the one-sided soft-thresholding operator [70]

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 majorization-minimization under our simplifying assumptions takes the form

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 ∇ Z R ( Z ) and ∇ 2 Z R ( Z ) : for every Z and every ∆ and every orthogonal D ,

Hence the quadratic Taylor expansion (79) can be written equivalently as

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 ∗ Z ℓ +1 / 2 ; moreover, because our approximation arguments above are not affected by left-multiplication of Z ℓ +1 / 2 by an orthogonal matrix (this operation does not change the norms of the columns of Z ℓ +1 / 2 , 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-minimization iteration is given by

for any orthogonal dictionary D . This gives an update quite similar to 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 a natural white-box version of this part of the architecture, along with the natural interpretation that its purpose is to sparsify the compressed tokens Z ℓ +1 / 2 in a (learnable) dictionary , which accords with recent empirical studies [79].

Other architectures? As we mentioned at the start of this section, the preceding derivation is performed in the most elementary possible setting in order to demonstrate the majorizationminimization 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:

Beyond exactly-incoherent features : our derivations above assumed that the incoming representations Z ℓ +1 / 2 were already maximal for the expansion term R in (13). It is desirable to obtain a 'perturbative' derivation, which applies in cases where Z ℓ +1 / 2 is not fully orthogonal, but instead near-orthogonal, 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 to involve additional normalization.

Beyond orthogonal dictionaries : The symmetries of the expansion term R in (13) may be followed to lead to a pair of dictionaries D and D ′ and an objective that sparsifies DZD ′ . This type of transformation is suggestive of popular architectures that mix over tokens [54, 67], however we consider the simpler form DZ 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 D which are complete and near-orthogonal. Adapting the derivation to overcomplete dictionaries 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 U [ K ] and their embedding dimensions.

A.3.1 Auxiliary Lemmas

Lemma 5. Consider the function

where α > 0 is a constant. Then we have

and the Hessian operator ∇ 2 Z R ( Z ): R d × N → R d × N satisfies that for any ∆ ∈ R d × N ,

Proof. The gradient calculation follows from [46], for example. For the Hessian, we use the usual approach to calculating derivatives: if ∆ is any matrix with the same shape as Z and t > 0 ,

valid since R is smooth. We have

where in the fourth line we require that t is sufficiently close to 0 in order to invoke the Neumann series. First, notice that the term involving ∆ ∗ ∆ does not play a role in the final expression: after we differentiate with respect to t and take a limit t → 0 , terms arising due to differentiation of t ↦→ t ∆ ∗ ∆ go to zero, because whenever the summation index k > 0 we have a term ( -αt ) k that goes to zero as t → 0 . We thus obtain with the product rule

Lemma 6. One has

Proof. Fix ∆ satisfying ∥ ∆ ∥ F ≤ 1 . By the triangle inequality,

For the first term, we note that

and since ( I + α Z ∗ t Z t ) -1 ⪯ I , we obtain from Cauchy-Schwarz 13

We can use a similar idea to control the second term. We have from the triangle inequality

For the first term, we have

The last estimate follows from a computation using the SVD of Z t . Meanwhile, we have for the second term by a similar argument (using the fact that the singular values of A and A ∗ are identical for any matrix A )

where once again the estimate follows from a computation involving the SVD of Z t (together with the fact that the function σ ↦→ σ/ (1 + ασ 2 ) is bounded on σ ≥ 0 by 1 / (2 √ α ) ). Putting it together, we have obtained

which gives the claim after taking suprema.

B Additional Experiments and Details

In this section, we provide details about our experiments, and report the results of additional experiments that were not covered in the main text. CRATE takes arguably the most basic design choices possible, and so we do not attempt to directly compete with state-of-the-art performance from heavily engineered and empirically designed transformers. The results of our experiments are meant to convey a few core messages:

Despite not being engineered to compete with the state-of-the-art, CRATE performs strongly on large-scale real-world datasets , including classification on ImageNet-1K. CRATE also achieves strong transfer learning performance.

Because our model is designed through unrolled optimization of a well-understood objective, each layer is interpretable . In particular, we can analyze the performance of CRATE, as well as design network modifications, on a layer-wise basis . This is powered by an arguably unparalleled level of insight into the role of each operator in our network.

We make the simplest possible choices during the design of CRATE , but these can be changed easily while keeping the same framework . We study a few modifications later in this section (Appendix B.4) and show that they do not significantly hurt empirical performance, but emphasize here that there is significant potential for improvement with different architecture choices (and in particular a different theoretical analysis).

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 to the conceptual architecture proposed in Sections 2.5 and 3, we make the following change for the sake of implementation simplicity:

In the compression step, replace the term p Nϵ 2 [ U 1 , . . . , U K ] in the MSSA operator with another trainable parameter W ∈ R d × pK . Thus the MSSA block becomes

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 × 10 -4 , weight decay as 0.5, and batch size as 2,048. We incorporate a warm-up strategy with a linear increase over 5 epochs, followed by training the models for a total of 150 epochs with cosine decay. For data augmentation, we only apply the standard techniques, random cropping and random horizontal flipping, on the ImageNet-1K dataset. We apply label smoothing with smoothing parameter 0 . 1 . One training epoch of CRATE -Base takes around 240 seconds using 16 A100 40GB GPUs.

Fine-tuning. We fine-tune our pre-trained CRATE and ViT models on the following target datasets: CIFAR10/CIFAR100 [10], Oxford Flowers-102 [7], Oxford-IIIT-Pets [16]. We also evaluate our pre-trained models on the commonly used ImageNet Real [36] benchmark. For each fine-tuning task, we use the AdamW optimizer [26]. We configure the learning rate as 5 × 10 -5 , weight decay as 0.01, and batch size to be 512. To allow transfer learning, we first resize our input data to 224. For data augmentations, we also adopt several standard techniques: random cropping, random horizontal flipping, and random augmentation (with number of transformations n = 2 and magnitude of transformations m = 14 ). 14

Algorithm 1: PyTorch-style pseudocode for CRATENetwork

Algorithm 2: Pytorch Style Pseudocode for Transformer Block in CRATE

Algorithm 3: Pseudocode for Attention and FeedForward

B.2 Experimental Results

In this subsection, we provide additional experimental results on CRATE, including layer-wise measurements, visualizations, as well as ablation studies.

B.2.1 Layer-wise Evaluation and Visualization

Layer-wise evaluation of compression and sparsity. Similar to Figure 3, we conduct the layerwise evaluation of compression term and sparsity for CRATE-Tiny, CRATE-Base, and CRATE-Large. We observe similar behavior as mentioned in Section 3.1: both the compression term and the sparsity term improves as the layer index increases.

2

4

6

8

10

12

Layer index -

400

450

500

550

600

650

700

R c

(

Z

) [SSA block]

Measure coding rate across layers

(a)

Compression (Model: CRATE-Tiny).

2

4

6

8

10

12

Layer index -

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

Sparsity [ISTA block]

Measure output sparsity across layers

(b)

Sparsity (Model: CRATE-Tiny).

2

4

6

8

10

12

Layer index -

600

800

1000

1200

1400

R c

(

Z

) [SSA block]

Measure coding rate across layers

(c)

Compression (Model: CRATE-Base).

2

4

6

8

10

12

Layer index -

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Sparsity [ISTA block]

Measure output sparsity across layers

(d)

Sparsity (Model: CRATE-Base).

0

5

10

15

20

25

Layer index -

1000

1200

1400

1600

1800

2000

R c

(

Z

) [SSA block]

Measure coding rate across layers

(e)

Compression (Model: CRATE-Large).

0

5

10

15

20

25

Layer index -

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Sparsity [ISTA block]

Measure output sparsity across layers

(f)

Sparsity (Model: CRATE-Large).

Visualizing layer-wise token representations. In Figure 6, we visualize the token representations Z ℓ at different layers ℓ ∈ { 1 , . . . , 12 } . We provide more results evaluated on other samples in Appendix B.2.2.

Visualizing layer-wise subspaces in multi-head self-attention. We provide the visualization of U ℓ [ K ] in Figure 7.

Layer =1

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(a)

ℓ

= 1

.

Layer =2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(b)

ℓ

= 2

.

Layer =3

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(c)

ℓ

= 3

.

Layer =4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(d)

ℓ

= 4

.

Layer =5

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(e)

ℓ

= 5

.

Layer =6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(f)

ℓ

= 6

.

Layer =7

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(g)

ℓ

= 7

.

Layer =8

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(h)

ℓ

= 8

.

Layer =9

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(i)

ℓ

= 9

.

Layer

=10

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(j)

ℓ

= 10

.

Layer

=11

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(k)

ℓ

= 11

.

Layer

=12

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

(l)

ℓ

= 12

.

Layer =1

0.2

0.4

0.6

0.8

1.0

(a)

ℓ

= 1

.

Layer =2

0.2

0.4

0.6

0.8

1.0

(b)

ℓ

= 2

.

Layer =3

0.2

0.4

0.6

0.8

1.0

(c)

ℓ

= 3

.

Layer =4

0.2

0.4

0.6

0.8

1.0

(d)

ℓ

= 4

.

Layer =5

0.2

0.4

0.6

0.8

1.0

(e)

ℓ

= 5

.

Layer =6

0.2

0.4

0.6

0.8

1.0

(f)

ℓ

= 6

.

Layer =7

0.2

0.4

0.6

0.8

1.0

(g)

ℓ

= 7

.

Layer =8

0.2

0.4

0.6

0.8

1.0

(h)

ℓ

= 8

.

Layer =9

0.2

0.4

0.6

0.8

1.0

(i)

ℓ

= 9

.

Layer

=10

0.2

0.4

0.6

0.8

1.0

(j)

ℓ

= 10

.

Layer

=11

0.2

0.4

0.6

0.8

1.0

(k)

ℓ

= 11

.

Layer

=12

0.2

0.4

0.6

0.8

1.0

(l)

ℓ

= 12

.

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

Layer =1

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =3

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =5

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =7

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =8

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =9

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=10

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=11

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=12

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =1

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =3

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =5

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =7

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =8

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =9

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=10

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=11

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=12

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =1

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =3

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =5

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =7

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =8

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =9

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=10

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=11

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=12

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =1

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =3

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =5

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =6

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =7

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =8

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer =9

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=10

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=11

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Layer

=12

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

B.3 CRATE Ablation

Hyperparameters of CRATE. In Table 2, we present evaluation of CRATE trained with various parameters. More specifically, we investigate the effect of number of epochs, weight decay, learning rate, step size ( η ) and the regularization term ( λ ) in ISTA block. As shown in Table 2, CRATE demonstrates consistently satisfactory performance across a diverse range of hyperparameters.

B.4 Exploring Architecture Variants

In this section, we explore the two following alternative architectures. One architecture involves a modification to the attention mechanism, while the other involves a modification to the sparsification mechanism. Again, we re-emphasize that these choices, although principled, are entirely modular and the choices we make here still lead to very simple architectures. A more sophisticated analysis may lead to different, more complicated architectures that perform better in practice. The architectures we experiment with are:

Compression-inspired attention mechanism: revert the change in (115). That is, the attention mechanism implements (11) and (12) directly.

Majorization-minimization proximal step sparsification: instead of (17), implement (92).

We obtain the following classification results in Table 3. After conducting additional simplifications to the network architecture (i.e., imposing additional constraints to the network architecture design), we discover that CRATE maintains reasonable performance on ImageNet-1K.

B.5 Sparse Coding vs. Non-Negative Sparse Coding

In the main body, we used a non-negative sparse coding formulation (16) to obtain the ISTA block as an unrolled proximal gradient step in (17). Suppose that we hadn't done this, and instead directly computed an ISTA block using an unrolled proximal gradient step on (15). Such a block would give the following update rule:

where S λη is the soft-thresholding function

applied element-wise to its input matrix. The resulting architecture would be an alternative to CRATE; below, we discuss some empirical and theoretical similarities and differences between the two formulations and architectures.

Empirical evaluation of soft-thresholding-based architecture. We summarize the results of CRATE with soft-thresholding activation in Table 4. We use λ = 10 in S λη and set all other hyperparameters the same as in the original CRATE-Base evaluation on ImageNet-1K. We find that such a soft-thresholding model achieves slightly worse performance-a drop of 3.2% top-1 accuracy-compared to the default CRATE-Base (with ReLU activation).

Potential theoretical justification for the performance differential. Previous work [49] studied the phase collapse mechanism for understanding the non-linearities and the convolutional filters used in CNNs such as ResNet [23] on classification tasks. Specifically, they found that replacing the phase collapses with thresholding operators which enforce sparsity largely degrades the classification performance of CNNs. The effect of phase collapse analyzed in [49] is to better separate out the means of different classes within a classification task. This may account in part for the increase in classification accuracy reported in Table 4. On the other hand, we believe that the CRATE architecture will be applicable beyond just classification tasks. In CRATE, the purpose of the training process is to learn the local signal models at each layer (see e.g., Section 2.5). From this perspective, so long as the downstream training task requires semantically meaningful representations of the data distribution, the exact training configuration is of secondary importance. In particular, we may use self-supervised learning methods such as (masked) autoencoding to learn the signal models, whence there may not be any well-defined notion of class mean. In such cases, a priori we may expect both soft thresholding and nonnegative soft thresholding to perform comparably well. We leave the verification of this to future work.

Theoretical justification for non-negative sparse coding. The sparse rate reduction formulation (1) does not include a non-negative constraint, and the token representations have marginal distribution equal to a mixture of zero-mean Gaussians (which are symmetric around 0 ). Below, we argue that the non-negative sparse rate reduction optimization and the regular sparse rate reduction optimization engender representations which are qualitatively similar in many ways, which confirms our conceptual understanding of how CRATE performs structured representation learning.

First, we formalize the non-negative sparse rate reduction optimization problem. Let χ be the characteristic function (with codomain { 0 , ∞} ) of its input proposition. Then the non-negative analogue to (1) is

Although formal analysis of the optimal points of the sparse rate reduction maximization problem (1) or its nonnegative variant (118) is out of scope of this work, we see that the rate reduction maximization (i.e., max f ∈F [∆ R ( Z ; U [ K ] )] has optimal points characterized similarly to [46, Theorem A.6], namely that the representation of each distribution in the mixture is supported on a subspace with nearly isotropic covariance on this subspace, and the supporting subspaces are (nearly) orthogonal. Adding the sparsity term λ ∥ Z ∥ 0 for some regularizer λ would enforce the axis-alignment of the supporting subspaces; when adding in addition the nonnegativity term χ ( Z ≥ 0) , following through the proof of [46, Theorem A.6] suggests that the argument goes through with suitable modifications (in particular, considering the conclusions for the covariance rather than ZZ ⊤ ). This sketch suggests that the statistical and geometric properties of the optimal representation remain the same when adding the non-negative constraint to the sparse rate reduction formulation. We leave a detailed proof to future work.

B.6 Pre-training on ImageNet-21K

We 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-1K. As

shown in Table 5, with the CRATE-Base model (22.80M parameters), we achieve 80.2% top-1 accuracy; this is comparable to ViT-Base ( 86M parameters, 83.9%) [40] with around 25% of the parameters. For pre-training on ImageNet-21K, we configure the learning rate to 1 × 10 -4 , set the weight decay to 0.05, and use a batch size of 4,096. The total number of epochs is 90, with 10 warmup epochs. For fine-tuning on ImageNet-1K, we use the same set of parameters as described in Appendix B.1.2, with the exception of setting the learning rate to 5 × 10 -5 and having a total of 50 epochs.

B.7 Evaluating R c and sparsity for ViT

We conduct experiments to evaluate the R c and sparsity of token representations from each layer of a pre-trained ViT-Base (downloaded from https://github.com/huggingface/ pytorch-image-models ). We summarize the results in Figure 12. We find that without our whitebox design, the vanilla ViT does not optimize our proposed sparse rate reduction objective. This contrasts with the results shown in Figures 3 and 4 of the work, wherein we can observe that the compression term R c and sparsity value decrease layerwise for CRATE, in accordance with our theory.

2

4

6

8

10

12

Layer index -

350

400

450

500

550

600

650

R c

(

Z

) [SSA block]

Measure coding rate across layers

ViT-S

2

4

6

8

10

12

Layer index -

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity [ISTA block]

Measure output sparsity across layers

ViT-S

2

4

6

8

10

12

Layer index -

0.75

0.80

0.85

0.90

0.95

1.00

Sparsity [ISTA block]

Measure output sparsity across layers

Threshold 1

Threshold 0.5

Threshold 0.1

2

4

6

8

10

12

Layer index -

600

700

800

900

1000

1100

R c

(

Z

) [SSA block]

Measure coding rate across layers

ViT-B

2

4

6

8

10

12

Layer index -

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity [ISTA block]

Measure output sparsity across layers

ViT-B

2

4

6

8

10

12

Layer index -

0.84

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

Sparsity [ISTA block]

Measure output sparsity across layers

Threshold 1

Threshold 0.5

Threshold 0.1

2

4

6

8

10

12

Layer index -

600

700

800

900

1000

1100

R c

(

Z

) [SSA block]

Measure coding rate across layers

CRATE-S

2

4

6

8

10

12

Layer index -

0.20

0.25

0.30

0.35

0.40

0.45

0.50

Sparsity [ISTA block]

Measure output sparsity across layers

CRATE-S

2

4

6

8

10

12

Layer index -

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

Sparsity [ISTA block]

Measure output sparsity across layers

Threshold 1

Threshold 0.5

Threshold 0.1

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOverall 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\nWEAKNESSES:\nThe 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?\n\nQUESTIONS:\nsee above.\n\nLIMITATIONS:\nsee above.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nUnderstanding 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.\n\nWEAKNESSES:\nAddressing 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\n\nQUESTIONS:\nWhat 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nTo 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.\n\nQUESTIONS:\nHow 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 ?\n\nLIMITATIONS:\nApart the aspect mentioned below, this is a really strong work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

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

STRENGTHS:
1) The paper is overall well written and easy to follow. Related Work section includes comprehensive surveys. Formulas are clearly explained. Experiments come with detailed settings.

2) The idea is novel and interesting. Especially, deriving multi-head attention from Maximizing Code Reduction is novel.

3) The proposed white-box architecture is verified on real-world datasets such as ImageNet-1k and compared to ViT models.

WEAKNESSES:
1) Although the paper claims that the proposed white-box model is competitive to ViT models, the numerical results seem not strong. On ImageNet-1k, the proposed model is clearly underperforming even with larger number of parameters.

2) The real power of ViTs is on high accuracy regime, where the model size is large. The authors only consider small model regime with low to medium accuracy, which lacks of convincing.

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

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10

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

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

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

QUESTIONS:
No.

LIMITATIONS:
No.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-12