Masked Completion via Structured Diffusion with White-Box Transformers
Authors: Druv Pai
Sam Buchanan
Ziyang Wu
Yaodong Yu
Yi Ma
Venue: ICLR 2024 poster (2024)
Keywords: masked autoencoding
matrix completion
structured representation learning
white-box transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We exploit a connection between denoising diffusion models, compression, and masked completion to construct white-box masked autoencoders from first principles.

Abstract: Modern learning frameworks often train deep neural networks with massive amounts of unlabeled data to learn representations by solving simple pretext tasks, then use the representations as foundations for downstream tasks. These networks are empirically designed; as such, they are usually not interpretable, their representations are not structured, and their designs are potentially redundant. White-box deep networks, in which each layer explicitly identifies and transforms structures in the data, present a promising alternative. However, existing white-box architectures have only been shown to work at scale in supervised settings with labeled data, such as classification. In this work, we provide the first instantiation of the white-box design paradigm that can be applied to large-scale unsupervised representation learning. We do this by exploiting a fundamental connection between diffusion, compression, and (masked) completion, deriving a deep transformer-like masked autoencoder architecture, called CRATE-MAE, in which the role of each layer is  mathematically fully interpretable:  they transform the data distribution to and from a structured representation. Extensive empirical evaluations confirm our analytical insights. CRATE-MAE demonstrates highly promising performance on large-scale imagery datasets while using only ~30% of the parameters compared to the standard masked autoencoder with the same model configuration. The representations learned by CRATE-MAE have explicit structure and also contain semantic meaning.

Published as a conference paper at ICLR 2024
MASKED COMPLETION VIA STRUCTURED DIFFUSION WITH WHITE-BOX TRANSFORMERS
Druv Pai" UC Berkeley
Ziyang Wu UC Berkeley
Sam Buchanan TTCC
Yaodong Yu UC Berkeley
Yi Ma
UC Berkeley & HKU
ABSTRACT
Modern learning frameworks often train deep neural networks with massive amounts of unlabeled data learn representations by solving simple pretext tasks, then use the representations as foundations for downstream tasks. These networks are empirically designed; as such, they are usually not interpretable, their repre- sentations are not structured, and their designs potentially redundant. White- box deep networks, in which each layer explicitly identifies and transforms struc- tures in the data, present promising alternative. However, existing white-box architectures have only been shown to work at scale in supervised settings with labeled data, such as classification. In this work, we provide the first instantia- tion of the white-box design paradigm that can be applied to large-scale unsuper- vised representation learning. We do this by exploiting fundamental connec- tion between diffusion, compression, and (masked) completion, deriving a deep transformer-like masked autoencoder architecture, called CR CRATE-MAF, which the role of each layer is mathematically fully interpretable: they transform the data distribution to and from a structured representation. Extensive empirical evalua- tions confirm our analytical insights. CRATE-MAE demonstrates highly promis- ing performance on large-scale imagery datasets while using only ~30 the parameters compared to the standard masked autoencoder with the same model configuration. The representations learned by CRATE-MAE have explicit structure and also contain semantic meaning- Code is available GitHub.
INTRODUCTION
In recent years, deep learning has been called upon to process continually larger quantities of high- dimensional, noisy, and unlabeled data, key property which makes these ever-larger tasks tractable is that the high-dimensional data tends to have low-dimensional geometric and statistical structure. Modern deep networks tend to learn (implicit or explicit) representations of this structure, which then used to efficiently perform downstream tasks. Learning these representations is thus of central importance machine learning, and there are so far several common methodologies for this task. We focus our attention below on approaches that incrementally transform the data towards the end representation with simple, alternatives is postponed to Appendix A.
hatpaaaticccly-inttteettttl primitives. Discussion of popular
Denoising-diffusse models for high-dimensional data. popular method for learning implicit representations of high-dimensional data is learning denoise: model that can denoise, i.e., remove noise from a corrupted observation from the data distribution (to the extent information- theoretically possible), can be chained across noise levels transform the data distribution and from certain highly structured distributions, such as an isotropic Gaussian, enabling efficient sam- pling (Ho al., 2020; Hyvärinen, 2005; Kadkhodaie 2015; Song et al., 2021; 2023; Vincent, 2011). Crucially, in the case of data with low-dimensional structure srrueture-includingg the highly nonlinear structure characteristic of natural images these models can be learned efficiently (Chen et al., 2023; Moitra & Risteski, 2020; Oko et al., 2023), and as result this framework has significant practical impact (Rombach et al., 2022). Despite this progress, these techniques have been largely limited use in generative modeling; a key reason is the un- structured nature of the final 'noisy" state of the diffusion process, which makes it challenging to control and interpret the model's learned implicit representation of the data.
Simoncelli, 2021; Sohl-Dickstein et al.,
"Correspondence Druv Pai, druvpai i @berkel ley. edu.

---

Published as a conference paper at ICLR 2024
White-box models and structured representation learning. In contrast, white-box models designed to produce explicit and structured representations of the data distribution according to a desired parsimonious configuration, such as sparsity (Gregor & LeCun, 2010; Zhai et al., 2020) or (piecewise) linearity (Chan et al., 2022). Recent work (Chan et al., 2022; Yu et al., 2023a) has built white-box deep networks via unrolled optimization: namely, obtain representations with desired set of properties, one constructs an objective function which encourages these desiderata, then constructs a deep network where each layer designed iteratively optimize the objective. This builds deep networks as chain of operators, representing well-understood optimization prim- itives, which iteratively transform the representations the desired structure. For example, Yu al. (2023a) uses an information -hhoorrtc objective promoting lossy compression of the data towards fixed statistical structure to build a transformer-like architecture named CRATE in the above manner. However, such-obtained deep networks have yet to be constructed for most unsupervised contexts. The fundamental difficulty here is that decoder networks must map from representations to data, and hence invert (in a distributional sense) the transformations made to the data distribution by the encoder. This renders the unrolled optimization approach used to construct white-box encoders such CRATE infeasible for constructing decoders, and instead demands a fine-grained understanding of the operators that implement the encoder, and their (distributional) inverses.
Our contributions. To overcome this difficulty and extend the applicability of white-box models to unsupervised settings, we demonstrate in this work that these two paradigms have more in com- mon than previously appreciated. First, we show quantitatively that under certain natural regimes, denoising and compression are highly similar primitive data processing operations: when the target distribution has low-dimensional structure, both operations implement a projection operation onto this structure. Second, using this insight, we demonstrate quantitative connection between unrolled discretized diffusion models and unrolled significant expansion of the existing conceptual toolkit for developing white-box neural network architectures, which we use to derive white-box transformer-like encoder and decoder architectures that together form autoencoding model that we call CRATE-MAE, illustrated in Fig. 1. We eval- uate CRATE-MAL on the challenging masked autoencoding task (He et al., 2022) and demonstrate promising performance with large parameter savings traditional masked autoencoders, along with many side benefits such as emergence of semantic meaning in the representations.
tewpiiiaattoononnstuuttttt deep networks. This leads to a
2 APPROACH
2.1 SETUP AND NOTATION
We use the same notation and basic problem setup in Yu et al. (2023a). Namely, we have some matrix-valued random variable X = [x]]...,,=N E DDNN representing the data, where the æ, e RD are called "tokens" and may be arbitrarily correlated. To obtain representations of the input, we learn an encoder f: I DxN + DdxN, N; our representations are denoted by the random variable Z = f(x) = [2 xN| E R N where the token representations are z; E Rd. In the autoencoding
setup, we also learn decoder g: RdxN + D such that X X [1 ........... xxl = +(().
Our encoder and decoder will be deep neural networks, and as such they will be composed of sev- eral, say T., layers each. Write { f o... f o fure and g = g gh- o o go, where f RdxNN" - RaxNN and g²: TR - RdxN are the gth layer the encoder and decoder respec- tively, and frre; IR DxN + ddd and gost __x.N + DDN are the pre- and post-processing lay- ers respectively. The input to the gth layer of the encoder is denoted Z' of |2f,...,, zNv E _dx and the inpul to the gth layer of the decoder is denoted Y = [yf... y RRdx
2.2 DESIDERATA, OBJECTIVE, AND OPTIMIZATION
Our goal is to use the encoder f and decoder g to learn representations Z which are parsimonious (Ma et al., 2022) and invertible; namely, they have low-dimensional, sparse, (piecewise) linear geo- metric and statistical structure, and are (approximately) bijective with the original data X. Yu et al. (2023a) proposes to implement these desiderata by positing signal model for the representations:
low-Dimensionll Gaussian Mixture Codebook. Let Z = [xı...., ØN] E dxN be a random matrix. We impose the following statistical model on Z, parameterized by orthonormal bases
2

---

Published as a conference paper at ICLR 2024
Naak X)
Z Y
Mask
Figure 1: Diagram of the overall white-box CRATE-MAE pipeline, illustrating the end-to-end (masked) autoencoding process. The token representations transformed iteratively towards parsimonious (e-g.. compressed and sparse) representation by cach encoder layer f Furthermore, such representations are trans- formed back to the original image by the decoder layers g². Each encoder layer f is meant to be (partially) inverted by corresponding decoder layer gh-
Uik|= = (Un)kE(K) E (RdxP)); K each token z; has marginal distribution given by
2; ai Vi [N]
(2.1)
where (8i)kE[N] E [K]N are random variables corresponding to the subspace indices, and (OG)EE(N E (BPYN are zero-mean Gaussian variables. If we optionally specify a noise parame- ter Z O, we mean that we "diffuse" the tokens with Gaussian noise: by an abuse of notation, each token zi has marginal distribution given by
2, =
U, ar ow0 E [N]
(2.2)
where (Wi)EEN E (RR)) are i.i.d. standard Gaussian variables, independent of and os. 8;
the Up are sufficiently incoherent and axis-aligned, we expect such representations to maximize the spurse rute reduction objective function (Yu et al., 2023a):
Ez Ez|AR(Z UKK] All [o] Ez (R(Z) (Z UKK) [|o]
(2.3)
where R and R are lossy coding rates, or rate distortions (Cover, 1999), which are estimates for the number of bits required to encode the sample up precision > 0 using a Gaussian codebook, both unconditionally (for R), and conditioned on the samples being drawn from Uy summed over all l (for R²). Closed-form form estimates (Ma et al., 2007: Yu et al., 2023a) for such rate distortions are:
a d
R(Z)= log det IN + +az =Z)
(2.4)
Ne²
( UK)
log det (Ix + 3(U 2) (U, Z)),
P N
(2.5)
Notably, R is a measure of compression against our statistical structure it measures how closely the overall distribution of tokens Z fit a Gaussian mixture on UμK-- Meanwhile, the other two terms R and || llo ensure non-collapse and sparsity of the representations, respectively.
Following Yu et al. (2023a), one then constructs a deep network that incrementally optimizes the sparse rate reduction order to transform the data distribution towards the desired parsimonious configuration (2.1). Specifically, Yu et al. (2023a) proposed to construct the deep neural network f as a two-step alternating optimization procedure which compresses the input against the (learned) local signal model U(K] at layer f, by taking a step of gradient descent on R | U(K)) and subsequently taking step of proximal gradient descent on LASSO objective (Tibshirani, 1996; Wright & Ma, 2022) to sparsify the data in a (learned) dictionary D E Rdxd,
2+1/2 z MSSA(Z Z
((k) Z' RV ( U(K)
(2.6)
z = ISTA 1112 Dd) ~ arg min
1
21 + +XIIZ1
(2.7)
ZzO
2
3

---

Published as a conference paper at ICLR 2024
Multi Head Subspace Self-Attention (M5
Sparse Coding Proximal Step
(1STA)
compression
sparsification
Z
Z L+1/2
Z 1/2
z
The compresiion-pparsffccttoo iteration implemented by each layer of CRATE, and each coder layer CRATE-MAE. The compression step, implemented by the MSSA operator, projects the tokens Z' towards the subspace model UiK] to form Z z The sparsification step, implemented by the ISTA operator, rotates the tokens in z =+1222 towards the coordinate axes, using the sparsifying dictionary D to get z zz+1 The steps are performed sequence and comprise a single of the CRATE-MAE encoder.
Figure 2:
where MSSA(-), the Multi-head Subspace Self-Attention block (Yu et al., 2023a), is defined as
[UT Z) softmax (U z) [ Z))
MSSA(Z ([K] = 2
U
Ukl
(2.8)
(U 2) softmax (oaaaax((U22 (U) (U T 2))
and ISTA(-), the Iterative Shrinkage- Thresholding Algorithm block (Yu et al., 2023a), is defined as ISTA(Z D) = ReLU(Z- n1 (DZ Z) n/11).
(2.9)
The MSSA block is exactly the same as a multi-head self-attention block in transformer, with the changes that the Qu/Ks/V, blocks are replaced by single matrix Us in each head k. The resulting layer f thus bears significant resemblance to a transformer-lii block, and so the CRATE model is white-box transformer- like architecture constructed via unrolled optimization. Such CRATE models obtain competitive performance on standard tasks while enjoying many side benefits (Yu et al., 2023a;b), yet they have so far only been trained for supervised learning. In the sequel, we introduce paradigm to obtain fully white-box networks for unsupervised learning, such as autoencoding, through a novel understanding of the CRATE model's s distributional layerwise inverse.
2.3 UNIFYING COMPRESSION AND DENOISING
To transform representations the idealized signal model given by (2.1), we seek to iteratively remove the disturbances or deviations of each sample from this signal model. One way to perform this task is to perform lossy data compression (Ma et al., 2007; Psenka et al., 2023; Yu et al., 2020; 2023a): compressed versions of the data, without ancillary disturhances, form the representations. This approach has been favored in the construction of previous white-box deep networks, such as CRATE described above, due to the existence of explicit informatiion-hhorree sion. In this case, the term RZ UIK]), defined in (2.5), measures the lossy compression of the representations Z against the class of statistical models given by (2.1). Thus, an operation to minimize R such as (2.6), implements step of compression learn better representations.
criteria for compres-
Another way to remove disturbances from the signal model (2.1), especially if the perturbed model has the noisy structure given in (2.2), is to denoise. When the data is highly structured or low- dimensional, one-step denoising becomes statistically and computationally difficult (Pedregosa, 2023). Hence the modern solution to this problem is via denoising diffusion models, which take many small denoising steps towards the data distribution at progressively decreasing noise levels (Ho et al., 2020; Karras et al., 2022; Song et al., 2021). Such models use estimates of the so-called score scorefuucttin function V log Pa (Hyvärinen, 2005), where pa is the probability density function of the noised input when the noise has standard deviation o > 0. At all sufficiently small values of a, the score function V log (Z) for a particular noised input Z points towards the closest point to Z on the data distribution support (Chen al., 2023; Lu et al., 2023; Yu al., 2023a), or more generally the modes of the true data distribution, which guides the denoising diffusion model to project Z onto the support of the data distribution and diffuse within this support."
A mathematical exposition of diffusion models may be found in Appendix A.I.
4

---

Published as a conference paper at ICLR 2024
2
U2
( Unchs Z Mpchy
VV log
2j
(2 &
Zj
+V (2)
U1
U
Figure 3: Compression and denoising against the low-dimensional Gaussian mixture token model (2.1) are equivalent. Left: effect of compression against the low-dimensional Gaussian mixture model tokens (2.1), i.c., taking gradient steps on the coding rate (- UjK]) or equivalently, using the MSSA(- U(K]) operator - which is shown in Theorem to be equivalent to projecting onto the U-k- Right: the effect of denoising against (2.1), i.e., taking gradient steps on the score function of the noisy model (2.2) small noise levels o, or equivalently small times t. Up to scaling factors (not pictured), these two operations and have similar geometric and statistical interpretations as a projection onto the support of the data distribution. This connection motivates our structured denoising- denoising-diffusion diffusion framework, as elaborated in Section 2.3.
equivalent,
In the context of (2.1) and (2.2), both denoising and compression operations conceptually remove additive disturbances from the data, visualized in Figure 3. In the following result, make this qualitative observation mathematically precise: we show that under a simplified version of the as a projection onto
signal model (2.1), taking gradient step on R2 a compression primitive, the local signal model U UK)), just as with the denoising primitive of taking a gradient step on log Pa-
Theorem (Informal version of Theorem 3 in Appendix A.2). Suppose Z follows the noisy Gaus- sian codebook model (2.2), with infinitesimal noise level of > O and subspace memberships s; distributed as i...d. categorical random variables on the set of subspace indices {1..... .. K}, inde- pendently of all other sources of randomness. Suppose in addition that the number of tokens N, the representation dimension d, the number of subspaces K, and the subspace dimensions p have rela- tive sizes matching those of practical transformer architectures including the CRATF-MMA encoder (specified in detail in Assumption 2). Then the negative compression gradient - V. ( Ujk) points from 2 to the nearest Ut-
Theorem establishes in a representative special case of the Gaussian codebook model (2.1) that at low noise levels, compression against the signal model (2.1) is equivalent to denoising against (2.1).
In the sequel, we use this connection to understand the MSSA operators of the CRATE-MAE encoder, derived in Section 2.2 from different perspective, as realizing an incremental transformation of the data distribution towards the signal model (2.1) via approximate denoising. This important property guarantees that a corresponding deterministic diffusion process-namely, the time reversal of the denoising process-impee by MSSA. Because these approximate denoising processes transform the data towards a parametric structure, we call them structured denoising-diffusion processes.
inverse operator for the compression operation implemented
2.4 CONSTRUCTING A DISTRIBUTIONAALY--WEEEEE
TRANSFORMER LAYER
In Section 2.1, we described a method to construct a white-bor box transformer-like encoder network via unrolled optimization meant to compress the data against learned geometric and statistical structures, say against a distribution of tokens where each token marginally distributed as Gaussian mixture supported UpK In Section 2.3, described in general terms approach that relates denoising and compression to yield a conceptually similar network using the formalism of diffusion models, this time trainable via autoencoding. In this section, we carry out this procedure concretely obtain an encoder and decoder layer with similarly interpretable operational characteristics.
To measure compression, use the Rd function defined in (2.5). By using standard (reverse-time time) diffusion process with a scaling of R as a drop-in replacement for the score (see Appendix A.3 for details), obtain that such denoising diffusion process may be described by the following
5

---

Published as a conference paper at ICLR 2024
stochastic differential equation (SDE) (Song al., 2021).
(z - -- (2() d(K)) de + V2dB(I).
E [0.7],
(2.10)
where (B(t))eE is a Brownian motion. As design choice, we wish to assert that our encoder and decoder ought to be deterministic, particular preferring that our encoder-decoder architecture achieves sample-wise autoencoding as opposed to distribution-wiss autoencoding or generation. Thus we need to construct some ordinary differential equation (ODE) which transports the input probability distribution in the same way (2.10). Such equation is readily obtained as the probability flow ODE (Song et al., 2021), which itself is commonly used for denoising and sampling (Lu et al., 2022; Song al., 2021; 2023) and has the form
(2 = -(T VR (2()) U(K))
AI [0,7]. 7].
(2.11)
-
In particular, the (t) generated by (2.10) and (2.11) have the same law. first-order discretization
(see Appendix A.3) of (2.11) with step size x obtains the iteration:
z 2 z MSSA
U(k)) ~ Z
N (
(k))
(2.12)
where MSSA(-) was detined in (2.8). Similar to Yu et al. (2023a), in order to optimize the sparse rate reduction of the features, and in particular to sparsify them, we instantiate learnable dictionary D E ddd and sparsify against it, obtaining
z IZTA(Z 11/2 D³),
(2.13)
where ISTA(.) was defined in (2.9). Thus, we obtain step iteration for the gth encoder layer fe, where Ze = (Z)))
2²1 = Zz + Zt (K)) z ISTA( 14122 D
(2.14)
This is the same layer in CRATE, whose conceptual behavior is illustrated in Figure 2. This equiv- alence stems from the fact that the diffusion probability flow (2.11) is conceptually and mechanically similar gradient flow the compression objective in certain regimes, and so it demonstrates useful conceptual connection between discretized diffusion and unrolled optimization as iteratively compressing or denoising the signal against the learned data structures.
Note that we parameterized a different local signal model U(k] and dictionary D at each layer, despite the continuous-timg time flows in (2.11) using only one (i.e., the final) local signal model. This is because the sparsification step (2.13) transforms the data distribution, and so we require a different local signal model at each layer to model the new (more sparse) data distribution; see Figure I for intuition on the iterative transformations. Also, having a different signal model each layer may allow for more efficient iterative linearization and compression of highly nonlinear structures.
Now that we have shown how the structured diffusion approach can recover the original CRATE architecture (Yu et al., 2023a) an encoder in our autoencoding problem, we use our new approach to construct a novel matching decoder. The time reversal of the ODE (2.11) is:
1 (Y U(K) dl,
d)
VI E [0, T..
(2.15)
21
in the sense that the Y (T - generated by (2.15) has the same law as the Z(t) generated by (2.11), assuming compatible initial conditions. A first-order discretization of (2.15) obtains the iteration:
r = 2+122 (+1/2
Mikg) (k))= y + ( 2+11/2 V(k),, (2.16)
where Vix] (Vf..... VR and each V RRdx are the bases of the subspaces "anti-compress"
against. In our work, we treat them different from the corresponding U because the dis- cretization of (2.11) and (5) is imperfect, and thus we should not expect a 1-1 correspondence between local signal models in the encoder and decoder. To invert the effect of a sparsifying TSTA step. which our mental model in Figure 2 portrays as a rotation of the subspace supports to a more incoherent configuration, we multiply by another learnable dictionary EL E Iadxd obtaining
Y 2+1/2 E
y y ++1/2
MS MwSAAYY 441/2 V
(2.17)
This constructs the (& 1) layer of of our decoder. In the implementation. add layer normal- izations to ensure that the features are roughly constant-size so that the above approximations hold. Figure 4 has graphical depiction of the encoder and decoder layers.

---

Published as a conference paper at ICLR 2024
Z
LayerNorm
MSSA
LayerNorm
ISTA
z
y
MSSA
LayerNorm
Linear
LayerNorm
Y
Figure 4: Diagram of each encoder layer (top) and decoder layer (bottom) in CRATE-MAE. Notice that the two layers are highly anti-parallel each constructed do the operations the other reverse order. That is, in the decoder layer 9¹. the ISTA block of fl- is partially inverted first using a linear layer, then the MSSA block of f/-t 7 reversed; this order unravels the transformation done in ff fb-
2.5 A COMPLETE WHITE-BoC TRANSFORMER-LIKS ARCHITECTURE FOR AUTOENCODING
As previously discussed, the encoder is the concatenation a preprocewiing map f PxN
R N which has learnable parameters Wpre E R and ET E and has the form:
p (X) = W X +
EP
(2.18)
and L transformer-lik layers f": R + Rax given by
(²(2") = ISTA Z + MSSA Z'
U(K) D*),
E [7]]
(2.19)
omitting normalization for simplicity. The decoder the concatenation of transformer-lik layers g RaXN dx given by
g = E MS V(K))
Vl [4] -
(2.20)
with postprocessing map gost R + R x which is learnable linear map Wpost E R Dxd;
goet (y y
(2.21)
A full diagram of the autoencoding procedure is given in Figure 1.
Our training procedure seeks to learn and represent the structures in the data distribution. For this, we use a pretext task that measures the degree to which these structures have been learned: masked autoencoding (He et al., 2022), which "masks out"" a large percentage of randomly selected tokens the input X and then attempts to reconstruct the whole image, measuring success by the resulting autoencoding performance. Conceptually, masked autoencoding can be seen nonlinear general- ization of the classical matrix completion task, which exploits oi-deemnsional structure to impute missing entries in incomplete data; classical matrix completion can be solved efficiently if and only the data have low-dimensional structure (A (Amelunxen al., 2014; Candès & Recht, 2009; Wright & Ma, 2022). The masked autoencoding loss writes
[MAE(J (J.)) /g) = [lll( (Mask((X)))
X
(2.22)
Further implementation details of this architecture are discussed in Appendices B.1 and B.2.
EMPIRICAL EVALUATIONS
In this section, we conduct experiments evaluate CRATE-MAE on real-world datasets and both supervised and unsupervised tasks. Similarly to Yu et al. (2023a), CRATE-MAC is built using sim- ple design choices that we do not claim are optimal. We also do not claim that our results are optimally engineered; in particular, we do not use the extreme amount of computational resources required to obtain state-of-theart performance using vision transformer-buce bucked masked autoen- coders (MAEs) (He et al., 2022). Our goals in this section are to verify that white-box masked autoencoding model CRATE-MAE has promising performance and learns semantically meaningful representations, and that each operator in CRATE-MAN MAE aligns with theoretical design. We pro- vide additional experimental details Appendices B.1 and B.2.
Network architecture and training configuration. We implement the encoder and decoder archi- tectures described in Section 2, with few changes detailed in Appendix B. We consider different model sizes of CRATE- MAF by varying the token dimension d, number of heads K, and number of
7

---

Published as a conference paper at ICLR 2024
Table 1: Model configurations for different sizes of CRATE-MAE, parameter counts, and comparisons to VIT-MAE models from Gandelsman et al. (2022); He (2022). We observe that CRATE-MAB M----MA--aa Base uses around 30% the parameters ViT-MAE-Basse and similar number parameters as ViT-MAE-Small
Model Configuration
d
K
N
CRATE MAF Parameters
VIT-MAE # Parameters
Small (-S)
12
576
12
196
25.4M
47.6M
Base (-B)
12
768
12
196
44.6M
143.8M
Measure coding rate across layers
Massure output sparsity across layers
Sparsity (ISTA block]
y LSSA block)
val
Vol
Layer
Layer
Figure 5: Left: The compression measure H 2+1/2
Uxx) different layers of encoder. Right:
sparsity measure z (d-N,, different layers of encoder. Measurements were collected from CRATE-MAE-Base averaged over 10000 randomly chosen ImageNet samples. We observe that the compression and sparsity improve consistently over each layer and through the whole network.
Masked
VIT-MAE
CRATE-MAE
Original
Masked
MAE
CRATE- MAE
Original
Figure 6: Autoencoding visualizations of CRATE-MAE-Race and ViT-MAE-Base (He al., 2022) with 75% patches masked. We observe that the reconstructions from CRATE-MAE-Band are on par with recon- structions from ViT-MAE-Bat despite using < 1/3 of the parameters.
layers L; such parameters will be kept the same for the encoder and decoder, which is contrary to He et al. (2022) but in line with our white-box derivation. Table displays the CRATE-MAE model con- figurations and number of parameters, and compares with equivalent ViT-MAE model sizes (Gan- delsman et al., 2022; He et al., 2022), showing that CRATE-MAE uses around 30% of the parameters of MAE with the same model configuration. We consider ImageNet-IK (Deng et al., 2009) as the main experimental setting for our architecture. We apply the Adam W (Loshchilov & Hutter, 2019) optimizer to train CRATE-MAE models for both pre-training and fine-tuning. When fine-tuning, we also use several commonly used downstream datasets: CIFAR10, CIFAR100 (Krizhevsky et al., 2009), Oxford Flowers (Nilsback & Zisserman, 2008), and Oxford-IIT-Pets (Parkhi et al., 2012).
Layer-wise function analysis. First, we confirm that our model actually does do layer-wise com- pression and sparsification, confirming our conceptual understanding as described in Section 2. In Figure 5, observe that each layer of the encoder tends compress and sparsify the input features, confirming our theoretical designing of the role of cach operator in the network.
Autoencoding performance. In Figure 6, we qualitatively compare the masked autoencoding per- formance of CRATE- BRATE-MAE-BEE Base to ViT-MAE- Base (He et al., 2022). We observe that both models able to reconstruct the data well, despite CRATE-MAN using less than third of the parameters of ViT-MAE. In Table 4 (deferred to Appendix B.4) we display the average reconstruction loss of MAFR-MAA-Base and ViT-MAF--asee showing similar quantitative conclusion.
8

---

Published as a conference paper at ICLR 2024
Table 2: Top-1 classification accuracy of CRATE-MAE models when pre-trained ImageNet-IK and evaluated via fine-tuning or linear probing various datasets. We compare CRATE-MAL standard ViT-MAE models with many more parameters. Our results show that CRATE-MAE achieves competitive per- formance on this transfer learning task when either fine-tuning the whole model or just the classification head.
Classification Accuracy
CRATE-MAE-S
CRATE-MAE-B
VIT-MAE-S
ViT-MAE-B
Fine-Tuning
CIFARIO
96.2
96.8
97.6
98.5
CIFARI00
79.0
80.3
83.0
87.0
Oxford 1lwwers-122
71.7
78.5
84.2
92.5
Oxford-IIIT-Ptts
73.7
76.7
81.7
90.3
Linear Probing
CIFARIO
79.4
80.9
79.9
87.9
CIFARION
56.6
62.3
68.0
Oxford Flowers-
57.7
61.8
66.8
66.4
Oxford -ird--III-PPee--
40.6
46.2
51.8
80.1
(a) Visualizing PCA of token representations.
(b) Visualizing selected attention head outputs.
Figure 7: Left: Visualizations of the alignment of each image's token representations with the top three principal components (in red, blue, and green respectively) all token representations of images the given class. Right: Visualizations hand-picked attention map across all attention heads in the layer of the CRATE-MAE encoder for each image. We observe in Figure 7a that the top three principal components are aligned with tokens from parts of the image that carry its semantics, and in Figure 7b that the attention maps correctly "attend to" (activate strongly on) exactly the parts of the image which semantically meaningful.
Representation learning and emerging semantic properties. In Table 2, display the perfor- mance of CRATE-MAE models when fine-tuned or linear probed for supervised classification (precise method in Appendix B.1) on variety of datasets. We observe that the classification accuracies of CRATE-MAE MAE models are competitive with much larger ViT-MAE models. Moreover, the learned representations of CRATE-MAE carry useful semantic content. By taking the alignment of the rep- resentations of each token with the top few principal components of the representations of tokens in each class (precise details in Appendix B.3), we observe in Figure 7 (left) that the representa- tions are linearized, and that the top few principal components carry semantic structure. In Figure (right), we observe that the attention heads the MSSA operator in CRATE-MAE capture the se- manties of the input images. These properties have previously been observed in white-box models trained with supervised cross-entropy losses (Yu et al., 2023b); our results demonstrate that they are consequences of the white-box architecture, rather than the loss function.
CONCLUSION
In this work, we uncover quantitative connection between denoising and compression, and use it to design conceptual framework for building white-box (mathematically interpretable) transformer- like deep neural networks which can learn using unsupervised pretext tasks. such as masked au- toencoding. We show that such models are more parameter-efficient over their empirically designed cousins, achieve promising performance large-scale real-world imagery datasets, and learn struc- tured representations that contain semantic meaning. This work demonstrates the potential and practicality of white-box networks derived from first principles for tasks outside supervised classi- fication. We thus believe that this work helps to bridge the theory and practice of deep learning, by unifying both the conceptual and technical level many previously separated approaches includ- ing, but not limited to, diffusion and denoising, compression and rate reduction, transformers, and (masked) autoencoding.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The authors would like acknowledge help from Tianzhe Chu in preparing the manuscript. Druv Pai acknowledges support from a UC Berkeley College of Engineering fellowship. Yaodong Yu acknowledges support from the joint Simons Foundation-NSF DMS grant #2031899 and AI Com- munity Mini Grant from Future of Life Institute. Yi Ma acknowledes support from the joint Simons Foundation-NST DMS grant #2031899, the ONR grant 3000014-2-1-21022 and the University of Hong Kong.
REFERENCES
Samira Abnar and Willem H. Zuidema. Quantifying Attention Flow in Transformers. In Dan Ju- rafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault (eds.), Proceedings of the 58th An- nual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pp. 4190 4197. Association for Computational Linguistics, 2020. doi: 10. 10.18663/V1/2002. ACI.-35. URL https: doi. 10tg//1..1653//1/222222222 acl-main. 385.
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic Interpolants: A Unity- ing Framework for Flows and Diffusions. CoRR, abs/2303.08897, 2023, doi: A0.. 2303.08797. URL https: /oii. /rg//c. 48550/arrka /ar.iaariiiivvv. 2303. 08797.
Dennis Amelunxen, Martin Lotz, Michael B McCoy, and Joel A Tropp. Living on the edge: phase transitions in convex programs with random data. Information and Inference: A Journal of the IMA, 3(3):224-294, June 2014. ISSN 2049-8764. doi: 10. 10.1093/imaiai/aa0005. URL https:
//academic. oup. com/ www..oo/mmmmaa.aaaaattteteee.e..
4/112203
Shir Amir, Yossi Gandelsman, Shai Bagon, and Tali Dekel. Deep ViT Features as Dense Visual De- scriptors. CoRR, abs/21 12.05814, 2021. URL http://wwwwwwwwaa//////////a///a
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Variance- haattee---narrinne--ovvaaaaaaa Regu- larization for Self-supervised Learning. In The Tenth International Conference on Learning Representationss, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
https: // http:s///peeeeiieeeeeeeeeeeeeeee ew. net/ forum?i dte.:idddxx6666222000
Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle. Greedy Layer- wise Training of Deep Networks. In Bernhard Schölkopf, John C. Platt, and Thomas Hofmann (eds.), Advances in Neural Information Processing Systems 19, Proceedings of the Twentieth Annual Conference on Neural Information Processing Systems, Vancouver, British Columbia, Canada, December 4-7, 2006, pp. 153-160. MIT Press, 2006. URL https: /pwpcee/oeeeeeeooooooooooo ngs www.ieurpp.... cccc// 2aper/20066 hash/
Abst centtractttmll... html.
Emmanuel J. Candès and Benjamin Recht. Exact Matrix Completion via Convex Optimization. Found. Comput. Math., 96):717-772, 2009. doi: 10. 00077/102088009-90455 URL https:
//doi. /0. 000.:/0///11100000. 000000909905....
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging Properties in Self-supervised Vision Transformers. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10- 17, 2021, pp. //doi. wal.orgg110 10. 1109/ 09/ 20110///C////// wwCV4C9222222 2021 00951.
IEEE, 2021. doi: 10.1 00.1109//CC 100951. URL https:
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. J. Mach. Learn. Res., 23:114:1-114:10: 2022. URL http:// jmlr waagg/apeers. v23 22113322203333 html.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score Approximation, Estimation and Distribution Recovery of Diffusion Models on Low-dimensional Data. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). International Conference on Machine Learning, ICML 2023. 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 4672-4712. PMLR, 2023. URL https::////......... pww.ceeoceeiioooooooooooooooo..e ngs. ml r. www.ss/e chen230. htmll.
10

---

Published as a conference paper at ICLR 2024
Ting Chen, Simon Korbliith, Mohammad Norouzi, and Geoffrey E. Hinton. A Simple Frame- work for Contrastive Learning of Visual Representations. In Proceedings of the 37th Inter- national Conference Machine Learming. ICML 2020, 13-18 July 2020, Virtual Event, vol- ume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020. URL chen200 html.
http:// http://pocceediiggee ngs .ml mm wtt.:sml...p..r.... press/rsss/e.. vv11 11199 9
Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. Deconstructing Denoising Diffusion Mod- els for Self-supervised Learning. CoRR, abs/2401 14404, 2024. doi: 10.48550/ARXIV.2001.
14404. URL https::///// /tt:///dooo wto.oooog/l...... 10. 48550/ (tt/s5/////viii arxiv. 2401. 14404.
Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision Transformers Need Registers. CoRR, abs/2309. 16588, 2023. doi: 1044850/ARXI.23000.666 //doi. wwt.orgg//// 10. 48550/555/ aat.i..iiiivv.... 2658.
URL https:
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009). 20-25 June 2009, Miami, Florida, USA, pp. 248-255, IEEE Computer Society, 2009. doi: 1110//CVP..00095526666 1109/CVPR. 2 9009.0000000000000000 5206848.
URL https::/t//.o....... doii /10/
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Associa- lion, 1062446) 2011.
Yossi Gandelsman, Yu Sun, Xinlei Chen, and Alexei A. Efros. Test-time Training with Masked Au- toencoders. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 De- cember 9, 2022, 2022. URL ht http::::pppppprrss..eep cc/ paper f www.ee/papeeee/e -bw.:t53---bbtaaccc weppstracc-oofffeeeeeteeee
html.
Karol Gregor and Yann LeCun. Learning Fast Approximations of Sparse Coding. In Johannes Fürnkranz and Thorsten Joachims (eds.), Proceedings of the 27th International Conference Machine Learning (ICML-10), June 21-24, 2010, Haifa, Israel, pp. 399-406. Omnipress, 2010. URL www.tttpp/..om....m.m i cml cc/ Ste//ooofeeeeeeeeeee 22022010 0 wtpppprss.apppooooompppppeeemeea 449. .df..
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked Autoencoders Are Scalable Vision Learners In IEEE/CVF Conference on Computer Vision and Patlern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 15979-15988. IEEE, 2022. doi: 10. 10.110//92R22288...02 20222.05....2.1555.. URL https: /tt./ciii..o org wtt.ooo.go.......111......... CVPR52 888.6888688 2 2022.2222 01553.
Geoffrey E Hinton and Richard Zemel. Autoencoders, Minimum Description Length and Helmholtz Free Energy. In J. Cowan, G. Tesauro, and J. Alspector (eds.), Advances in Neural Information Processing Systems, volume 6. Morgan-Kaufmamn neurips wwtppppr_pepecceeeeeeee wet.ffeeee/ppe////eee 1993. URL http:9/eccccc4eeccfl8llla05333 663ee 29aef www.aaaa-eapeee....a.eeeaaaaeeee .ddf.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances Neural Information Processing Systems 33: Annual Con- ference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: 4c5bc
neurips.... cc wttpscpcpppp//////aaaaaaa html.
ab 1a11
Aapo Hyvärinen. Estimation of Non-normalized Statistical Models by Score Matching. J. Mach. Learn. Res., 6:695-709, 2005. URL http:// setp://jmir... hwt.mm..rrg//pppeae hatgg//apes///ee /v6 wwwyarinennnnssee. html.
11

---

Published as a conference paper at ICLR 2024
Zahra Kadkhodaie and Eero P. Simoncelli. Stochastic Solutions for Linear Inverse Prob- lems using the Prior Implicit in a Denoiser. In Mare Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Newral Information Processing Systems 34: Annual Conference on Neural Informa- fion Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 13242- 13254, 2021. URL https:// htte::://pprcceednng (e18999996ceede77e8eeee..
Abstract html.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the Design Space of Diffusion-based Generative Models. In Sanmi Koyejo, S. Mohamed, A. Agar- wal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Informa- tion Processing Systems 35: Anmual Conference on Neural Information Processing Sys- tems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http:/ /wpppappssss./ae wwwpers.niis.oo wnw.snips........... cttpppc.c/pppr__eeee Seseaper_flleepppeee htt.all//ppeer/22///aaaaaa a9884 6e9dcco0 cf wact
html.
Valentin Khrulkov, Gleb Ryzhakov, Andrei Chertkov, and Ivan V. Oseledets. Understanding DDPM Latent Codes Through Optimal Transport. In The Eleventh International Conference on Learning Representations. ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
URL https: / wwpp//eeeeeeeeeeeeeeeeeee net /
6IIhAx1144i.
Diederik P. Kingma and Max Welling. Auto-encoding Variational Bayes. Yoshua Bengio and
Yann LeCun (eds.), 2nd International Conference Learming Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL http:
//arxiv.org/ abs /tlbbs/1312 1111/11111100 6114.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The Efficient Transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https: /apenrevee/pereieeww net/ forum?id VKgNKkHtvB.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Mufan (Bill) Li, Mihai Nica, and Daniel M. Roy. The Neural Covariance SDE: Shaped Infinite Depth-and- Networks at Initialization. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http:/ http:://pppp....eee wwwappss.niss.... wttppssspppppppppppppeeppppeeeee wttpeeer.fisssssppeeeeeeeeee 4
htmml.
Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. net./ forum? dew.kgiiicccgoe.
OpenReview.net, 2019. URL
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Ad- vances in Neural Information Processing Systems 35: Annual Conference on Neural Informa- tion Processing Systems 2022, NewrlPS 2022, New Orleans, LA, USA, November 28 December http:sswssssspp.p.epeeeeeeeppspe cc/paper files/ waales/paper//e paaerr/222 //assh////////// html.
9, 2022, 2022. URL http: http://aaapaaaae/ 26Cal4acce2a8 9dad3 waweeaaaaaaaacdeeeeeeeeeeeeeeeee oat.ra/eeceeeetttttaatooooo-eee
Yubin Lu, Zhongjian Wang, and Guillaune Bal. Mathematical analysis of singularities in the diffu- sion model under the submanifold assumption, 2023.
Yi Ma. Harm Derksen, Wei Hong, and John Wright. Segmentation of Multivariate Mixed Data via Lossy Data Coding and Compression. IEEE Trans. Pattern Anal. Mach. Intell., 1949):1546 1562, 2007. doi: 10. 10.1109/PAAI..007.0855 URL https: /ttp//i.ooo////////ooo wogg/1..........o............... 1109 /MMM.. 2007. 1085.
12

---

Published as a conference paper at ICLR 2024
Yi Ma, Doris Tsao, and Heung- Yeung Shum. On the principles of Parsimony and Self-consistency for the emergence of intelligence. Frontiers Inf. Technol. Electron. Eng., 23(9) - 2022. doi: 10.1631/FITEE..20000 2200297. URL https: httpt::://do..ogg////// wtt./////......1/11.1 PWEEE.TTEEE.. TTT.TTEEEE.220999 2200297.
Song Mei and Yuchen Wu. Deep Networks Denoising Algorithms: Sample-efficient Learning of
Diffusion Models in High- High-dimensional Graphical Models. CoRR, abs/2309. 2023. doi: 10.
48550/ARXIV 1420. URL https: /tt.//dii... org 10. 48550 /tt//aaiivv.//..///v/ 230.2, 11420.
Annie Millet, David Nualart, and Marta Sanz. Integration by parts and time reversal for diffusion processes. The Annals of Probability, pp. 208-238, 1989.
Ankur Moitra and Andrej Risteski. Fast Convergence for Langevin Diffusion with Manifold Struc- ture. February 2020. URL http://::/aaaat org/ abs 2002 05576.
Maria-Elena Nilsback and Andrew Zisserman. Automated Flower Classification over Large Num- ber of Classes. In Sixth Indian Conference on Computer Vision, Graphics & Image Processing, ICVGIP 2008, Bhubaneswar, India, 16-19 December 2008, pp. 722-729. IEEE Computer Soci- cty, 2008. doi: 10.1100/CCVIPP22000..A8 URL https: dei.ciioooo http://.........111////1111 1109//111111110 ICVGIP. 2008. 47.
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion Models are Minimax Optimal Distri- bution Estimators. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 26517-26582. PMLR, 2023. URL https: http://rooceeeeeeeeooooooooooooc aige mmlr, press/////// v202/oko23a html.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Mahmoud Assran, Shang-Wen
Nicolas Ballas, Wojeiech Galuba, Russell Howes, Po-Yao Yao Huang, Shang- Wen Li, Ishan Misra, Michael G. Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jégou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning Robust Visual Features without Supervision. CoRR, abs/2304.07193, 2023. doi: 10.48500/ARX1.2304007999 URL fttps::/doi... http./////......///r........ www.rxiii......... 07193.
Omkar M. Parkhi, Andrea Vedaldi, Andrew Zisserman, and C. V. Jawahar. Cats and dogs. In 2012 IEEE Conference on Computer Vision and Pattern Recognition, Providence, RI, USA, June 16-21, 2012, pp. 3498-3505, IEEE Computer Society, 2012. doi: 10.1100/CVPR.221222245888 https: /ttpss//doll. www.diiooooo..... 1109 / CVWR.CCPRC// 2012. 6248092.
URL
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine Learning in Python. Journal of Machine Learning Research, 12:2825-288 2830,, 2011.
Fabian Pedregosa. On the Convergence of the Unadjusted Langevin Algorithm. 2023, URL https: //fa. //fa.biame net beet/bllo//222/ 2003//222///u 1t//0////llaae ulag/.
Michael Psenka, Druv Pai, Vishal Raman, Shankar Sastry, and Yi Ma. Representation Learning via Manifold Flattening and Reconstruction. CoRR, abs/2305.01777, 2023. doi: 10.48550/ARXIV. 2305.01777. URI. https: //oi org/ 48550/ 20..50//.......5...... 2305.01777
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic Backpropagation and Approximate Inference Deep Generative Models. In Proceedings of the 31th International
Conference on Machine Learning, ICML 2014, Beijing, China, 21-26 June 2014, volume 32 of JMLR Workshop and Conference Proceedings, pp. 1278-1286. JMLR.org, 2014. URL http:
//proceedings .lr www.m....ess......... www.pprsssseee.. /v32/ rttp/deeeeeeeeeee html.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution Image Synthesis with Latent Diffusion Models. In JEEE/CVF Conference on Com- puter Vision and Pattern Recognition, CVPR 2022, New Orleans. LA. USA. June 18-24, 2022, pp. 10674-10685. IEEE, 2022. doi: 1LL.11CVVS228..2222222222 2022. 01 042.
URL https: / doi.
hTg//10.100/CCPPPB
13

---

Published as a conference paper at ICLR 2024
Simo Särkkä and Amo Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep Unsu- pervised Learning using Nonequilibrium Thermodynamics. In Francis Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning ICML 2015, Lille, Frunce, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceed- ings, pp. 2256-2265. JML..org, 2015. URL http: /ttpc//dooeeeeeoeeoooooooc wttpooceeddiissm....r........ m.s.......res//111111 sohl-diccsteii155 htmll.
Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Flo- rence d' d'Alché-Bue, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Infor- mation Processing Systems 32: Annual Conference on Neural Information Processing Sys- tems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11895- 11907. 2019. URL https://prrceeedings neurips...... htcp//appr///1////// 20ah//sh/ .tml.
3001 500.eee277776657711cccc611113332
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based Generative Modeling through Stochastic Differential Equations. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL nttps://openreviies net forum?in PXTIG12RRES
Yang Song, Prafulla Dhariwal, Mark Chen, and lya Sutskever. Consistency Models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scar- lett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Hon- olulu, Hawaii, USA. volume 202 of Proceedings of Machine Learning Research, pp. 32211- 32252. PMLR, 2023. html.
URL https: https::/ http///pooeeeenggeeeeeeeeces maliiingss....... mlr wwt.:nsssmr.....e.....ee.e /v202/22222222 song23a.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267-22) 1996.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Informaionn Theory Workshop, ITW 2015, Jerusalem, Israel, April 26 May 1, 2015, pp. 1-5. IEEE, 2015. doi: 10.1009/T..2005571331669 URL https:://cc...ooeeet
1109/
ITW. 2015. 7133169..
Roman Vershynin. High- dineessimnnion probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018
Pascal Vincent. A Connection Between Score Matching and Denoising Autoencoders. Neural Comput., 23(7):1661-1674, 2011. doi: (.1162/NECOOAA),001144 URL https://doi.
org/10. 1 62 LBE//ONN_OOOO11_01 a__1__12_22200
John Wright and Yi Ma. High-dimensional data analysis with low-dimensional models: Principles, computation, and applications. Cambridge University Press, 2022.
Weilai Xiang. Hongyu Yang, Di Huang, and Yunhong Wang. Denoising Diffusson Autoen- coders are Unified Self-supervised Learnrr. In IEEE/CVE International Conference Com- puter Vision, ICCV 2023, Paris, France, October 1-6, 2023, pp. 15756-15766. IEEE, 2023. doi: D0.109//CCVV10002222230000 2023. 01448.
URL https:://oo. http:/doo.ooogrr/// wtw.////000 1109 ICCV51070.
Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. Learning Di- verse and Discriminative Representations via the Principle of Maximal Coding Rate Reduc- tion. In Hugo Larochelle. Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Con- ference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL htt httppp:://ppooeeeingg..eeeeee 6ad4 1 174ebal 9ecb5fed1 741 A ract html.
14

---

Published as a conference paper at ICLR 2024
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D. Haeffele, and Yi Ma. White-box Transformers via Sparse Rate Reduction. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023a. URL http: wwwppspppssppppppppap nips. nips. caciips.cpppppee saaiper_filessppeeeeee hww.es/paper////3/////aa Abstract
- Conference html.
Yaodong Yu, Tianzhe Chu, Shenghang Tong. Ziyang Wu, Druv Pai, Sam Buchanan, and Yi Ma. Emergence of Segmentation with Minimalistic White-box Transformers. CoRR, 1121/238..167,,, 2023b. doi: (044550AAXXIV23308..62) URL https: w1t.://.00/../0////..... 2308. 16271.
Yuexiang Zhai, Zitong Yang. Zhenyu Liao, John Wright,, and Yi Ma. Complete dictionary learn- ing via 4-norm maximization over the orthogonal group. The Journal of Machine Learning Research, 21()):6622-6689 2020.
15

---

Published as a conference paper at ICLR 2024
A OTHER RELATED WORK
In Section 1, we described the approaches to unsupervised learning of low-dimensional structures in the data that were most relevant the rest of the work. Here, we discuss some other popular alternatives for completeness.
Black-box unsupervised representation learning. On the other end from white-box models, which learn representations from data that have a priori desired structures, are black-box unsuper- vised learning methods which Iearn fully data-driven representations. One implementation of this principle includes contrastive learning, which learns representations from computing the statistics of multiple augmentations of the same data point (Bardes et al., 2022; Chen et al., 2020). Another an- gle is to seek representation with desirable characteristies for a specific task, such as classification; prior works have considered diffusion models as "representation learners" from this angle (Chen et al., 2024; Xiang et al., 2023). The notion of representation learning work, namely the transformation of the data distribution towards a structured form that preserves its essential information content, is different from the notion in this latter group of works. Still another implementation is that of autoencoding models, the most recently popular of which is the masked autoencoder (MAE) (He et al., 2022). Autoencoders attempt to build low-dimensional representa- tions of the data and use them to reconstruct input data (Hinton & Zemel, 1993; Kingma & Welling, 2014; Rezende et al., 2014; Tishby& & Zaslavsky, 2015); masked autoencoders specifically mask the input data in training and attempt to impute the missing entries through reconstruction.
interested in in this
The common point in all such unsupervised learning methods so far is that they use black-box neural networks, such as ResNets (Chen et al., 2020) or transformers (Caron et al., 2021), as their back-end, Thus, although they sometimes develop semantically meaningful representations of the data (Bardes et al., 2022; Caron et al., 2021; Chen et al., 2020), they are uninterpretable, and their training procedures and internal mechanisms are opaque.
Deep networks and stochastic dynamies. There are many quantitative rapprochements of deep learning and stochastic dynamics. The most well-known of these is diffusion models, which can be modeled as discretizations of Itö diffusion processes (Song et al., 2021). The neural network is usually trained to estimate the so-called score function. Diffusion models can be thought of as implementing a particular approximation to optimal transport between template distribution and the true data distribution (Khrulkov et al., 2023). Different types of stochastic dynamics useful for generative modeling may be derived from optimal transport between the data distribution and pre- specified template (Albergo et al., 2023; De Bortoli et al., 2021). However, diffusion models are unique among these methods in that they have an iterative denoising interpretation (Karras et al., 2022), which this work draws on. Such an interpretation has previously been used to construct deep denoising networks from unrolled diffusion processes (Mei & Wu, 2023), instead of just using the deep networks to do black-box estimation of the score function. Similar studies have interpreted deep networks as discretizations of diffusion processes without this particular denoising interpreta- tion (Li et al., 2022), but the aforementioned unrolled iterative denoising strategy is what we draw upon in this work.
Other related work. Here also diseuss some related work with regards to different modifica- tions of the transformer architecture and training procedures which interface well with our white-box design. For example, Kitaev et al. (2020) suggests that sharing the Q and K matrices in the regular transformer is mechanism to make the transformer more efficient at no performance cost. This choice is heuristic, whereas our white-box design suggests that Q, K, and V should be set equal, and as we see in the paper this comes with some small tradeoffs. Also, since white-box models are derived such that each layer has defined and understood role, is natural ask if such models can be trained layer-wise, i.e., one layer at a time (Bengio et al., 2006). While this is also possible, we leave it to future work; our experiments show that with end-to-end training, a vastly more common method to train deep networks, each layer still follows the role it was designed for.
A.1 AN OVERVIEW OF DIFFUSION PROCESSES
In this section, we give an overview of the basics of time-reversible Itô diffusion processes, the mathematical foundation for diffusion models. This is to make this paper more self-contained by
16

---

Published as a conference paper at ICLR 2024
providing knowledge about general diffusion processes that will apply our special models. The coverage adapts that of Karras et al. (2022); Millet et al. (1989); Song et al. (2021).
Consider generic Itō diffusion process (E))))) where x(t) is 12" random variable, given by the SDE
(z(t) (t). ()) dw)(t)
E(z(t),
z z(0) ~
VE [0
(A.1)
where is Brownian motion and P is some probability measure I (in this case representing the data distribution). Here the drift coefficient b: I2'h x R + IR' and diffusion coefficient E: IR x I + x are functions. To make sense of (A.1) and also verify the existence of strong (i.e., pathwise well-defined) solutions, we need some regularity on them, and we choose the following assumption:
A1. and E have some spatial smoothness and do not grow too fast, i.e., there is constant K 2 0 such that for all æ; E R we have
sup [| [xx, 1) E(Z, ()l [|b 1) b(2, 112) s K|| a E||2
(A.2)
xE(T.T]
sup tE[o,T]
[1 t)ll= + [16 1a & K + ||
(A.3)
In general, z(t) may not have density w.r.t. the Lebesgue measure on R For example, suppose that P is supported on some low-dimensional linear subspace (or even Dirac delta measure), and take Σ be the orthoprojector onto this subspace. Then xt) will be supported on this subspace for all / and thus not have a density w.r.t the Lebesgue measure. Thus, when further discussing processes of the type (A.1), we make the following assumption
A2. z(l) has a probability density function p(:,t) for all /. 0. This is guaranteed by either of the following conditions (Millet et al., 1989):
A2.1 b and Σ differentiable in (ar, and have Hölder-continuous derivatives, and P has density w.r.t. the Lebesgue measure;
A2.2 The event
(rank(E(=(s)),))) s)) m for all some neighborhood of 0}
(A.4)
happens P-almost surely.
Define v: RT R + mmxm by
(xx = M(x t))(
(A.5)
To discuss rine-reversibility we also need the following local integrability condition, which is an- other measure of sharp growth of the coefficients (or precisely their derivatives):
A3.
The functions (æ,t) + Va (V t))(xx t)) are integrable sets of the form D [to, for to > and D a bounded measurable subset of "!:
[ Iv= (V(a, t)p(x, t) dx dt < oc.
[''
(A.6)
To write the notation out more explicitly,
Væ ( (zx ()((x))))
V (Vi (((x,)))(xx1)
(A.7)
V (wm (x, t)p(a) t))
where
V ())t)(( t)) =
/ (a, ()(xxt))
(A.8)
dz,
where Vi is the 4th row of V transposed to a column, and t the (i. entry of w. Note that Millet al. (1989) phrases this in terms of local integrability condition each
17

---

Published as a conference paper at ICLR 2024
Ve (V t)plz, t)),, which would naturally give local integrability condition on Ve (2, t)p(a, lo) However, all norms on RT are equivalent, and so this leads to a local inte- grability condition for VV (v t)) t))ll2 produced. Note that the assumptions do not guarantee that the involved derivatives exist, in which case they are taken in the distributional (e.g.. weak) sense, whence they should exist (Millet et al., 1989).
Under assumptions A1-A3, Millet et al. (1989) guarantees the existence of another process tejor)) such that the laws of z() and E(T - t) are the same for all [0, This pro- cess (0)7)) is called the time reversal of (z (2)))) and is shown to have law given by
(t) b- (z)), t) dt + + (z(), dw (t), (0) p(,
V [0, (A.9)
where w² (() is an independent Brownian motion and
b- (æ, = -b(x, - t) + V a t)p(2,, -t]] bb(e T -
(A.10)
pl - 1) + vx Va,T + 1) 1)],
1) -
(A.11)
2" (x, = T'=
(A.12)
We would next like to develop an ODE which transports the probability mass P in the same way as (A.1) - namely, find another process (Z(t)) which has deterministic dynamics, yet has the same law as (=(t))ee(o,7) Song et al. (2021) looks at the Pekker-Plance equations (which can be defined, at least weak sense, under assumptions Al-A2) and manipulates them to get the following dynamics for (t):
(1) b(z(1), 1) dt,
z(0) ~ P.
WL [0,7],
(A.13)
1 Vn [V(x Op(e,, ()]
where
b(x, f) bbxr,t)
(A.14)
2
p(x,t)
1
1
b(x,, t)
2 V (xx 1 (xx
t)).
(A.15)
2
Now get similar process for z(t), namely process (zt))ee which evolves deterministically yet has the same law as (E(t)) we may either take the time reversal of (A.13) or apply the Fokker-Planck method to (A.9), in both cases obtaining the same dynamics:
dē((t) = b (z) dr,
5(0~>((.1)
Vt e [0,
(A.16)
where
5 (2,t)) = -b(x, T t))
(A.17)
V. [V T - t)p(2x - t)]
= -b(xz T -t) + 2
(A.18)
p(x, T t)
= bbb(x, t) + 1 v (z t) + ( a log p(x T t)].
(A.19)
The quantity v logp(æ, 1) is of central importance; denoted the score time f, and we use the notation s(x; t) = V. log p(x, t) for it. With this substitution, we have the following dynamics for our four processes:
dz(t)= (t dw =(0) ~ P dž(t)= = -()(((())TT)++ V. -t)
(A.20)
dt
(A.21)
+ =t) (t), ž ~P(TT)
(A.22)
dz(()= b(E(t), 1 - V. (((())) 1
((((()
dt, æ(0) P
(A.23)
1 Vo ((
dž(t)
b(((t,T + 2
(A.24)
-) + 22 - (5())
1) dt, ž(() T):
(A.25)
18

---

Published as a conference paper at ICLR 2024
In practice, one fits an estimator for s(:,)) and estimates p(-, T) and runs discretization of either (A.9) or (A.16) to sample approximately from P. One common instantiation used in diffusion models (Karras et al., 2022) is the so-called variance-exploding diffusion process, which has the coefficient settings
b(x, = [x,, = 22
(A.26)
which implies that
(xx t) = 21.
(A.27)
This means that the four specified processes are of the form
dx(() d2w()) =00) P
(A.28)
dm(t) = s( - t) dt + d2dw
( ~P(-1))
(A.29)
dz(t) sz((t), f) dt, 2(0) P
(A.30)
dB(()) (()(
{0) ~ p(,
(A.31)
Notice that the determinstic flows are actually gradient flows the score, which concretely reveals connection between sampling and optimization, and thus between diffusion models (precisely those which use the probability flow ODE to sample) and unrolled optimization networks.
In this context, we can also establish the connection between diffusion networks and iterative de- noising. In the vuriance-explodin setting, we have
<t) ~ (22(0)) 2tT),
(A.32)
which can be easily computed using results from, e.g.. Särkkä & Solin (2019). Thus z(t) is noisy version of z0), with noise level increasing monotonically with t, and sampling (0) from =(t) conceptually removes this noise. Concretely, Tweedie's 's formula (Efron, 2011) says that the optimal
denoising function [[2( xt)) has a simple form in terms of the score function:
([z(0) x(t)] xtt + 2t (z),,,, t).
(A.33)
In other words, the score function s points from the current iterate z(t) to the value of the optimal denoising function, so it is a negative multiple of the conditionally-expeeeed noise. Following the score by (stochastic) gradient flow or its discretization is thus equivalent iterative denoising.
A.2 COMPANION TO SECTION 2.3
In this section, we prove a formal version of the result Theorem I stated in Section 2.3. That is, we examine a basic yet representative instantiation of the signal model (2.2), and show that under this model, in natural regime of parameter scales motivated by the architecture of CRATE- MAE applied to standard image classification benchmarks, the operation implemented by taking a gradient step on the compression term of the sparse rate reduction objective (2.3) corresponds to projection operation at quantization scales g2 proportional to the size of the deviation. This leads us in particular to a formal version of the result Theorem 1.
Signal model. We consider an instantiation of the model (2.2), elaborated here. That is, we fix a distribution over tokens Z E RaxN induced by the following natural signal model: each token z, is drawn independently from the normalized isotropic Gaussian measure on one of K p-dimensional subspaces with orthonormal bases U1.... UK E dx which comprise the low-dimensional structure in the observed tokens, then corrupted with i.i.d. Gaussian noise A(, o² I); the subspace each token is drawn from is selected uniformly at random, independently of all other randomness the problem. This signal model therefore corresponds to the setting of uncorrelated tokens, with maximum entropy coordinate distributions within subspaces. is natural to first develop our theoret- ical understanding of the connection between compression and the score function in the uncorrelated setting, although in general, the ability of CRATE-MAE capture correlations in the data through the MSSA block is essential. In connection with the latter issue, we note that our proofs will gen- eralize straightforwardly the setting of "well-dispersed" correlated tokens: the discussion in Remark 5.
We make the further following assumptions within this model:
MMore precisely, 2; is distributed according to the pushforward of the normalized isotropic Gaussian mea- sure (0, 1 IT) on IR" by the bases Ug.
19

---

Published as a conference paper at ICLR 2024
1. Inspired by ablation in Yu et al. (2023a), which suggests that the learned PRATE-MAL model on supervised classification on ImageNet has signal models UR which are near- incoherent, we will assume that the subspaces Ug have pairwise orthogonal column spaces. Our proofs will generalize straightforwardly to the setting where the subspaces are merely incoherent: the discussion in Remark
2. We assume that the relative scales of these parameters conform the CRATE-MAE settings, trained on ImageNet: from Table 1, these parameters are
(a) d = 768;
(b) N 196;
(c) K 12;
(d) = d/K K = 64.
In particular, d > > and Kp = d.
These precise parameter values will not play a role in our analysis. We merely require the following quantitative relationships between the parameter values, which are general than the above precise settings.
Assumption 2. We have s S 1, U UK. = L for all + W, and the following parameter settings and scales:
d
Kp= d;
CVVNlog N < //NK, where C₁ the same as the universal constant C1 in the statement of Proposition 12;
6C2 < d, where C2 is the same the universal constant C3 in the statement of Propo- sition 15;
N < d, where C4 is the same as the universal constant C1 in Proposition 11;
Note: there is no self-reference, as the third inequality is not used to prove Proposition 12, the fourth used prove Proposition 15, and the fifth is not used to prove Proposition 11.
The first and second inequalities together imply in particular that 2 N/K. The third inequality implies that CIVN log < N/K. K. The first, second, and and third inequalities together imply that C CVN N log N, and that < N/K CIvN N log N < N/K < N/K + CIVN log N < N.
These inequalities are verifiable in practice
wishes to explicitly compute the absolute constants C1, C2, C3, C4. and indeed they hold for our PRATE-MAE-A -Bse model.
Formally, let (KK, p. o²) denote the probability measure on dR corresponding the noisy Gaus- sian mixture distribution specified above. We let Z₃ ~ denote an observation distributed according this signal model: formally, there exists a (random) map + 8;, for E [N] and s; E [K]. such that
Zei = U. ;4; + 8,, I= I,..,,,
(A.34)
where A 88
8N Lid. NO, o² I), and (independently) a miid No0, 1 IT. It is p
convenient to write this observation model in block form. To this end, let K₁ = E 11 1,,= for E [K] denote the number of times the k-th subspace is represented amongst the columns of Z₁ (a random variable). Then by rotational invariance of the Gaussian distribution, we have
Z U1A1 ... UKAK II A,
(A.35)
where denotes equality in distribution, E NNNN a uniformly random permutation matrix, and each A₁ E Roxke We also define X, to be the noise-free version of Z,-
Because of this equality in distribution, we will commit the mild abuse of notation of identifying the block representation (A.35) with the observation model (A.34) that follows the distribution μ.
20

---

Published as a conference paper at ICLR 2024
Denoising in the uncorrelated tokens model. In the uncorrelated tokens model (A.35), the marginal distribution each column of Z is identical, and equal to an equiproportinal mixture of (normalized) isotropic Gaussians the subspaces UI.... UR. convolved with the noise distribution (O, T) This marginal distribution was studied in Yu et al. (2023a), where it was shown that when the perturbation level o² + O, the score function for this marginal distribution approximately implements a projection operation onto the nearest subspace Uu.
Hence, we can connect compression, as implemented in the MSSA block of the CRATE-MAL archi- tecture, to denoising in the uncorrelated tokens model by showing that at similar local scales, and for suitable settings of the model parameters, the compression operation implements projection onto the wwwdimensiooal structure of the distribution, as well.
Compression operation. The MSSA block of the CRATE-MAL architecture arises from taking an (approximate) gradient step on the Ro term of the sparse rate reduction objective (2.3). This term writes
( URK)
Σ log det (I B(U Z) U Z)
(A.36)
where
P
B
(A.37)
Ne
and >> is the quantization error. Calculating the gradient, we have
K
V (z U(K)
(A.38)
k=1
Minimizing the sparse rate reduction objective corresponds to taking a gradient descent step on (- U(K]). Performing this operation the observation from the uncorrelated tokens model Zn. the output can be written as
Z = (R, UnK).
(A.39)
where 7 > 0 the step size.
Main result projection. We will see shortly that the behavior of the compression output (A.39) depends on the relative seales of the perturbation about the low-dimensional structure o² and the target quantization error s².
Theorem 3. There are universal constants C1, C2, C3, C4 such that the following holds. Sup- pose Assumption 2 holds, and moreover suppose that a S 1 and C & 1. Then with probability 3.
at least KC2 Czd CCNNK + N-2), it holds
1+ 1+8-1 X:
zz (A (A- 1P (SAII II
(A.40)
1 3-1
<CyKn (0 + to //d) + VK30 VN/d) + N/d
(A.41)
Here, Pūpe implements projection onto the relevant subspaces for each token in the limiting case as E - 0, and is precisely defined in (A.116) and (A.117).
We give the proof of Theorem 3 below. First, we make three remarks on interpreting the result, our technical assumptions, and our analysis.
Remark 4. Theorem 3 admits the following interesting interpretation in an asymptotic setting, where identify the leading-order behavior of the gradient and confirm our hypothesis about the connection between compression and score-following. Choose = 3 so that the guarantee in Theorem incurs some cancellation, and moreover delineate more precise dependencies on the RHS
21

---

Published as a conference paper at ICLR 2024
of the guarantee:
z+
(A- (AII I
1 1-8-1 X;
(A.42)
NK?²
N2K2g1 +o(1+ /N/d) +
o²d
NVK (1 /N/d)+ (1+///)) VN/d)
do²
(A.43)
d
2 SKOR + o²d NK² 2 d
0
d E,
(A.44)
where we used Assumption 2, which implies p = d/ K and N/d < We will check due course whether we have satisfied the hypotheses of Theorem 3, that this guarantee indeed applies. To this end, we optimize this bound as a function & > since this parameter of the compression model. The optimal is straightforward to compute using calculus: it satisfies
o2d N
K
N d
(A.45)
d
od
(A.46)
NK
a
VX
and the value of the residual arising from Theorem 3 with this choice of is no larger than an absolute constant multiple of
K²o2d N
No d
N
3/2
N
+
d
= Ko VKo+ + V
(A.47)
Moreover, with this choice of e, 3 satisfies
N2NK
8
(A.48)
N d
In particular, the condition 30 $ in Theorem demands
N SL
V
(A.49)
which holds sufficiently small a and sufficiently large d N, showing that Theorem 3 can be nontrivially applied in this setting. If consider simplifying limiting regime where N, d - too such that N/d + 0 and N/K- + +oo, we observe the following asymptotic behavior of the guarantee of Theorem 3:
z x (A-P Puna (AII I
x.
| S
& Ko0/2
(Ko)
(A.50)
1+
This demonstrates that gradient step R. performs denoising: there is leise-levelldeeendeet shrinkage effect applied to the signal X,, which vanishes as a + 0, and meanwhile the noise term A is reduced.
Moreover, as a - 0, we can express the limiting form of Puiky exactly as an orthogonal projection, since this drives 3 - 0: following (A.116) and (A.117), we have here
Puk
[P1
Pk],
(A.51)
where
Pk Σ Up projim(A U..
(A.52)
k' k'
This shows that, in an asymptotic sense, a gradient step on R serves to suppress the effect of the perturbation applied to the observations Zy about the local signal model X,. This verifies claim
22

---

Published as a conference paper at ICLR 2024
previously that in this setting, there is correspondence between score-following algorithm and compression-bast based approach: locally, both project the observations onto the structures of the signal model.
can be shown moreover that the shrinkage effect on X; demonstrated here appears as a conse- quence of using the R "compression" term for the gradient step in CRATE-MAE; when the gradient step is taken instead on the full A R rate reduction objective (which is computationally prohibitive, of course), there is zero shrinkage, and perfect denoising performed for wider variety of step sizes than the choice made here. We see the introduction of this shrinkage effect this as the price of constructing an efficient and interpretable network architecture. In practice, the ISTA block of CRATE-MAE counteracts this shrinkage effect, which is anyways minor at reasonable parameter scales.
Remark 5. We have made two assumptions which may not hold exactly in practice: namely, we have assumed that the Us 's have orthogonal columns, namely U Up = k== I, and we have assumed that the linear combination coefficients As that form the matrix X; are i.i.d. samples from Gaussian distributions. Both these assumptions can be made more realistic, at the cost of additional (non-instructive) complexity in the analysis; we briefly go over how.
Relaxing the orthogonality condition U U2 = 1 to near-orthogonality. namely | U2- I| for small v, observed in practice (Yu et al., 2023a) would introduce additional small error terms in the proof, say polynomial v. The magnitudes of these errors could in principle be precisely tracked, whence one could obtain similar result to Theorem 3.
Secondly, we have assumed that the A 's have independent columns which are sampled from (the same) Gaussian distribution. However, in the conceptual framework for CRATE-MAE, the joint distribution (and in particular the correlations) between the tokens in order to obtain good performance for our model. Our analysis is not completely agnostic this fact; as we will see, the proof of Theorem only leverages the independence of the columns of each A 's in order to obtain gego-prbbability upper bounds on the smallest and largest singular value of the token matrices. If these bounds were ensured by some other method, such as appropriate normalization and incoherence, similar conclusion to Theorem 3 could hold in the more realistic correlated tokens model. Going beyond well-conditioned token matrices for each subspace would require additional modeling assumptions, and additional investigative experimental work to determine realistic basis for such assumptions.
exploit
Remark 6. We have not attempted to optimize constants or rates of concentration in the proof of Theorem 3, preferring instead to pursue a straightforward analysis that leads to qualitative inter- pretation of the behavior of the rate reduction gradient our model problem. Minor improvements to the concentration analysis would enable the parameter scaling requirements in Assumption 2 to be relaxed slightly, and the probability bound in Theorem 3 that scales as K/N2 can easily be improved to any positive power of 1/N.
Proof of Theorem 3. We start by noticing that, by orthonormality of the subspaces Us, we have by (A.35)
U Z
0
A
0 + U A
(A.53)
so that
1-1
8
1I
Zx) TU 2.)
U
IIT
T
8 I + A.A.
+E II.
D
(A.54)
23

---

Published as a conference paper at ICLR 2024
because permutation matrices orthogonal matrices, and where the perturbation 3₁ is defined by
A
E IIA UUU AII +
A A
U.A
A, A
U Ak K
A
(A.55)
and where we have defined (implicitly) in addition
[
A = AnI
(A.56)
The matrix D > 0, So we can write
( + (U, 2,)
U 2x)
DE (I+ II,
(A.57)
from which it follows
U Ze 1 (U
I
),) U Z
(A.58)
= =(0
(B I A. A&) 0 U AI D (I+E
II. (A.59)
The task before us is therefore control Dx < in order to apply the Neumann series to further simplify this expression. We will do this in stages: first, we invoke several auxiliary lemmas construct a nigh-probability event which the random quantities in the preceding expression controlled about their nominal values; next, we show that the Neumann series can be applied on this event and main term extracted; finally, we simplify this main term further in order establish the claimed expression.
High--proabiility event construction. In order achieve the appropriate control on all random quantities, we would like to construct a high hiig-prrobbbiity event on which the random quantities are not too large. By Propositions 9, 10 and 11 and union bound, there exist universal constants C, for which
| < (C1 + NNdd) A C
IP
Vk E [K]: Ak E
2 1 - C,K(e Cod +e- -CN/KK + 2 (A.60)
[
A I|I S C3/N/d
The event we compute the probability of, which we denote F², , precisely the good event that we want. Formally,
A S a /N/d) A <1+ || A. - I| C. CayN/a
Vk [K] Vk [K]
(A.61)
We know that Ex occurs with high probability, and are able to strongly control the random quantities to the degree desired.
Main term extraction. By Lemma 8 and our hypotheses on the problem parameters, we have on Ex that
E |1 < Cão < 1.
(A.62)
We can therefore apply the Neumann series obtain
U Z (8- (U, Z,) U. Z,)
(A.63)
=((o
- U, AlI " (1 (I-E, (-." II.
(A.64)
Again on E*, we have
Σ 1y (S.DE'))
E '1' '1p < C(Bo)²
1- C3o S C3
(A.65)
1=2
24

---

Published as a conference paper at ICLR 2024
Moreover, as in the proof of Lemma 8, we have the previous event that
U AlI D
C3o..
(A.66)
Thus, we define a "main term"
M [ Ax Al Ax)
0] (I -E UN
ATT D II,
(A.67)
we have on the same event as previously
| Z (8-1 (U, 2 U 2,) - M1
<C(B))
(A.68)
To conclude, need only study this main term, since Up has operator norm
Simplifying the main term Our approach will be to control the main term M, around a simpler expression, using basic perturbation theory; by the triangle inequality for the operator norm, this will give control of the desired gradient term. After distributing, Mg a sum of three terms; we will start with the simplest term. We first compute
UR AlI D UU (A
Ds (B-LI+ AT Ab)
3AK
(A.69)
We are going to argue that the residual
| AI D
U [3
B
(A.70)
small. To this end, note that by the fact that Ux has unit operator norm
|0
AII
D
U [3A,
A K
(A.71)
<
A (8- AA)
(A.72)
A (8
I+ A. A&) I|
(A.73)
- A
A&)
(A.74)
By (A.61) and (A.120) from Lemma 7, the second term here is controlled on E For the first term, note that by definition and the fact that the unit sphere is invariant to rotations (and permutations are rotations),
A
sup
Aull,
sup |||A|
ull2
(A.75)
||
[ulla<1
sup |ulla$
u
(A.76)
where u, are coordinate--ubsst--ndudddd partitions of the vector u induced by those of AIl This yields immediately
sup 1ulla11
A u S
max A KE
sup
u lla
VK
max
A+
(A.77)
kE(K]
by the triangle inequality and inequalities for (?) nonms. Similarly, choosing specific in the operator norm expression, namely one that is supported entirely on one of the coordinate partitions ui shows that
|
A 11ll2
(A.78)
for any é, whence
(A.79)
K
follows that we control the first term above on E* Combining this reasoning, we conclude from the above
| AIT DE"
U AA
BAk
(A.80)
C'/N/d
So + VN/d)
(A.81)
1+3-1
1+ 8+1
So(1 So(1+C/N//)
(A.82)
25

---

Published as a conference paper at ICLR 2024
where the second line uses Assumption 2 remove the higher-order residual.
Next, we recall that E is a sum of two terms; we will do one term at a time for concision. We have first
0] IT
U
AIT
(A.83)
4
A
= =[0 1I AT
0
U
(A.84)
A
A
= Ax I Am)
|
Ak]
(A.85)
We then multiply this term by Dj on the right to get the term that appears in Mj (ignoring the multiplication on the right by II. because does not change operator norms). In particular, we will control
I AK
D
(A.86)
showing that this term small. We will accomplish this with the block diagonal structure of D indeed, this gives that D is obtained by blockwise inversion, and hence
A (3 I
A R
A
U
[1
DI
(A.87)
Ax ATA) AUU BA
A A8 A Ab)
BAK]]
]|| (A.88)
S V K max ||A6(8- A A
- AUUUT (8- A "1,
(A.89)
(5 1 A. AN) A 11. RYR
(A.90)
where the second line uses (A.77). We will give a coarse control of this term term-the error could be improved further by exploiting more thoroughly independence of the blocks A to show that the maximum over k' * k the last line of the preceding expression is smaller. We have by submuliiplicatiityy of the operator norm and the triangle inequality
Ax(B A AT AR)
AUU Ax(3-1 " AR)
(A.91)
2
CVNTU
1 1-1 +
IAUURU
(A.92)
1+3-1
< (1+C UAUU
(A.93)
where the first line uses Lemma 7, and the second line uses Assumption to simplify the residual as above. We have meanwhile from the definition of E
a U Ax s [|Alll (1+VN
(A.94)
because U an orthogonal projection, and again using Assumption 2 to simplify the residual. We can argue analogously to simplify the other term in the maximum appearing above, and this yields
Ax A A - A S V330 (1+C/N/a) (1+ /N/d).
AK] D
(A.95)
(A.96)
where we used the fact that 8 < and the rest of Assumption 2 implies that 8 > 1. This residual simplifies using Assumption 2 to
AA - A
D VK30² (I+C/N/a)
(A.97)
26

---

Published as a conference paper at ICLR 2024
Next, examine the last term, which is the other summand arising in the definition of E have
We
A.U.A.
0
0] U. A
A + A
A AK
AUA
(A.98)
Ax 1 A2) [ATUTA A
U AK]
(A.99)
Now multiplying on the right by D gives the term (again ignoring the rightt-mutiplicatioon by II, which does not affect the operator norm)
(18-1I+AAAA |SATU ... (ATURAR AT A.) (3-11 ALA
BATU, Ax
(A.100)
We will argue that this term is close to the term
A II A A.) [ A
0
BA AK]
(A.101)
in operator norm. The argument is similar the preceding arguments: for this, it suffices to bound
[ which is the same as controlling the operator norm of the nonzero block. Using submultiplicativity and Lemma along with the simplifications we have done above leveraging Assumption 2, we obtain
Ax I 1 Av) (ATU,A+ A. Aw) (8
o// (A.102)
A + A A U.U + ATU, AA) (8- ATA)) A[-I+A[[AAT
(A.103)
$ (1+C/N77) ( + A.U, A
(A.104)
Meanwhile, on E* we have the operator norm of A and A& controlled, using again (A.79). Ap- plying then the triangle inequality and submultiplicativity. we obtain
A _ A + A A S a (1 VN/d).
(A.105)
again simplifying with Assumption 2. This shows that (A.100) is close to (A.101) with deviations of the order $ 501 /N/d).
Aggregating the previous results. Combining our perturbation analysis above, we have estab- lished control
M
(I- AM I A A.) AT)
U [A
BAK A
(A.106)
+ 0 ...
AT
0
n
(A.107)
5a(1+ N/d) + VK d/d).
(A.108)
It is convenient to include one additional stage of simplification here: namely, we use Lemma once more simplify the second term in the nominal value of Mj appearing here. Namely, have (arguing as we have above, once again)
AB I + A A+)
143-1 A
(A.109)
1
A A (8-I++ ATA
-1
(A.110)
1+
3-1
/
N/d,
(A.111)
27

---

Published as a conference paper at ICLR 2024
from which it follows
Ma - (T Ax(3 A Ax) A U [AA
AAk]
(A.112)
+ 0
1+8 A
(A.113)
So(1 /N/d) + VL N/d) VN/d.
(A.114)
Mcanwhile, recall the residual of scale - (a3)² arising when we controlled the gradient term around Mş:
| z (3-1 (U 2 U -
- -M
C(B0)²
(A.115)
Combining these two bounds with the triangle inequality controls the gradient term around its nom- inal value. Now, we sum these errors over k (again with the triangle inequality) to obtain control of the aggregate gradient around its nominal value. We introduce notation to concisely capture the accumulations of the (approximate) orthogonal projections arising in the nominal value of the main term: for each [K], define
Pk EU
Un (I- Av Ap) UN K
(A.116)
and define overall (approximate) projection operator (which block matrices partitioned compatibly with the class sizes NR, as in (A.35)) by
Punn
[Pi
Px]
(A.117)
Then the above argument implies
1
V R Rd Zn U(K) - Punng (B (BAII
1 8 X;
(A.118)
$
+0(1 to
/N/d) + VK30 (N//// + /NNd)
(A.119)
which is enough to conclude.
A.2.1 KEY AUXILIARY LEMMAS
In this section we state and prove two key concentration inequalities that are used the proof of the main theorem. They rely on simpler results which will be conveyed in subsequent subsections.
Lemma 7. There exist universal constants C, C' such that the following holds. Let d. P, N, K N be such that Assumption holds. Let Ag. k E [K] be defined as above. Let E be the good event defined in (A.61). If E* occurs, then for K [K] we have
18-11+ (8-1++ Ax)
1
CVN (1+ B-1)
(A.120)
1
and in addition
C'VN/d (1+
Ax 'I AT AT)
A
5
(A.121)
1 1+
Proof. Since E. holds, for all & E [K] we have
| < 1 Ci C.VN// N/d,
|A] A 1|| S N2/
(A.122)
By Assumption we have || A A /||
1, so A. Ag is well-conditioned. Write
= A A I
(A.123)
28

---

Published as a conference paper at ICLR 2024
so that A1
-1+++-11++ +
(A.124)
1
I
1+ 8-1 3
(A.125)
1+
1
(A.126)
1+3-1 1=
1 + 3-1
1
3
(A.127)
1+ 1+
8-1
1 8-1
1 8-1
==1
by the Neumann series. This gives us
(8 I+ Ax)
1
Σ
86
E
(A.128)
1+
1+ 3-1
1 8-1
j=1
1
E
(A.129)
1+ 3-1
1+ 3-1
j=1
C2VN/d
(A.130)
1+ 8-1
1+
Ca C2VNId
(A.131)
1++B-1)(1 + 8-1 -021-C2///
Meanwhile, by Assumption 2, it holds
CavN/d V166
(A.132)
so it follows
C2VN/d + 8- C2 N/d
$
(A.133)
By the submultiplicatiiity of the operator norm, we thus have
1
Ax(8-114 A. Aa)
1+ Av
1 Axll (8- AT)
(A.134)
1+ 3-1
[1+C
$ (1+ C2/ N/d)
(A.135)
2 + Ciy N/dC2vNij
(A.136)
1+
C2 N// +/dd//CCC... C.CN/d
=
(A.137)
1+
By Assumption 2, we have that there exists some absolute constant C3 > 0 with C3- < /N/d, which gives
3-1 Av S (G++000,,,
Ax(8-'I4 An)
(A.138)
1+ 1+3-1
as desired.
Lemma 8. There exist universal constants C1,C₂ > such that the following holds. Let d. d.p, N, K N be such that Assumption 2 holds. Let As. k [K], be defined as above. Let Dy be defined as (A.54). Let Ey be defined as in (A.55). Let E* be the good event defined in (A.61). If E occurs, then for lo E [K] we have
E D -1 C1 +o( VN/d)).
(A.139)
29

---

Published as a conference paper at ICLR 2024
Proof. Since we have
1-
Dx
B I + A
(A.140)
3 B-1
it holds that
731
D
(3-1 + AT)
(A.141)
81_
We will use the straightforward estimate D. S | and bound the two matrices' operator norms individually. By the previous expression,
|
max 11 A Ax)
- < 3,
(A.142)
because A A > 0, need only control the operator norm of E To this end, note the convenient expression
E = II U.U AIT + 28ym ((AnT)T
[0
U
0))
(A.143)
where sym(-) denotes the symmetric part operator. By the triangle incquality, the operator norm in the previous expression.
of E larger than the sum of the operator norms of each The operator norm of the first term is no larger than || because II is permutation matrix and UU, is an orthogonal projection. Meanwhile, using that the symmetric part operator is the orthogonal projection onto the space of symmetric matrices, it follows
2aym ( 0
U
0)] - 2/ (AII") T [0
ULA
Olll
(A.144)
and then we find as above that the RHS is larger than |||||| A& Since the good event E defined in (A.61) holds by assumption, we have that there are constants C1, C2 > 0 such that
N
| S
C1+
C1
(A.145)
N | 51+ Cay
(A.146)
By Assumption 2 have d 2 N. so that V /N/d S Therefore we have
| S (C₁ + 1) = Cza
(A.147)
for C33 = C1 + another universal constant. Thus on this good event we have
2 A ||I A S Cza (1+ (1+C2/N/6)
(A.148)
Therefore, have
+ B || < | 2 A
(A.149)
$ Cf
(A.150)
< C +0 C2/N/d)
(A.151)
where C₁ max{Ca CCs, C3} is another universal constant. Thus Ex we have
D < C4 o(1 Ca/N/d) s Cs,8l02 + +o(1+ VN/d))
(A.152)
for Css 0 another obvious universal constant.
30

---

Published as a conference paper at ICLR 2024
A.2.2 CONCENTRATION INEQUALITIES FOR OUR SETTING
In this section prove some simple concentration inequalities that are adapted to the problem setting. These results are used to prove the key lemmata above, and indeed are also invoked in the main theorem. They follow from even simpler concentration inequalities that are abstracted away from the problem setting, which we discuss in the following subsections.
Proposition 9. There are universal constants Ci C2 C3 > such that the following holds. Let d, .N E N be such that Assumption 2 holds. Let A E R N be defined as above. Then
N (00200(+))))) (++
P P P
I Cze Czd
(A.153)
Proof. We use Proposition 14 with the parameters q d. n N, and .//o/d which obtains
svd/o VN - 11) A&A vd
P||||| s) C exp
d
NNN++C C2Vd) (A.154)
C2vd
Notice that we have
svd/o VN
8
N
N (VN + +2VV)) =0 ( + C
-1
1,
C2vd
C2
0
Vd
(A.155)
To make the squared term equal to 1, we pick
== σ (1 d + 2C2
N
+202
(A.156)
which gives
20' + ) Cze
|
(A.157)
Proposition 10. There universal constants C1 C2 C3, C1 >> 0 such that the following holds. Let N, K E N be such that Assumption holds. Let As, h E [K], be defined as above. Then
N d < e2ex
N
C4
PP | > + CiV
-Css
-
+
(A.158)
K
N
Proof. By Propositions 12 and 14 with parameters n = n, k: define
K, q = p. and = 1/ VD, if we
N
N K
Nmin
Ci CI
Nlog
Maxx
+ +I Nlg
(A.159)
then we have
Nanax
C2
P||lA||
P
> s) S
P[[|AA|
s K2 =npP n]
(A.160)
N²
Namin
N Nana
Czexp
-11
1/D VP
-
P P[Ka = n +
C2
(A.161)
CIV4
N2
N=
for all obeying
8 x /D ( (p + Cav C.V.a Nunx
1
(A.162)
Nmax
= 1 Viv
(A.)63)
p
31

---

Published as a conference paper at ICLR 2024
Thus have that the concentration holds for all obeying
Mmax
2 1 + Cap
(A.164)
P
In order to cancel out the most interior terms, we choose
Nmax
8 ==1 ==1-201p
(A.165)
p
This choice obtains
|||| A& > 8) < Σ
Niaxx
Czexp
-n
SVD vD
-1]
PlK = m) +
C2
(A.166)
Civn
N2
n=
2
Nmav
Nmax
C2
Σ
exp
-1 2
-1
P(K₂ = n] +
(A.167)
N
#= Nuin
21
21
Nmax
C2
Czexp(-n)P(Ks = n] +
(A.168)
N/2
Nuin
Nmax
C2
Mocpp-NNin)PPf = n] +
(A.169)
N2
Nmin
C2
C3exp( Nmin)
(A.170)
N2
N
C2
C3
+ C1 No N
(A.171)
K
N2
N
C2
S C3 exp
+
2
V
N Vog N
N2
(A.172)
1 N
C2
S C3exp exp
(A.173)
K
N2
To obtain the conclusion of the theorem, note that any such that
Nmax = 1 1 +24
N/K + + Volog = 1+ +020
N
VN
VVlog log N
11 + 2Cs
Ci
p
d
(A.174)
enjoys the same high-probability bound. By Assumption we have
N C
VNlog. N
N
1
N/K
1 2C,1
31+
(A.175)
d
d
2
P
N
N d
3
N
=1 =1+20
1 +1112
(A.176)
d
whence the ultimate conclusion is obtained by combining constants.
Proposition 11. There are universal constants C1, C2 Cs, C4 0 such that the following holds. Letp, N, E N be such that Assumption 2 holds. Let Ap, k E [K], be defined as above. Then
N
C4
P
A Av
1||
<
C2 exp
-C3
+
(A.177)
K
N
32

---

Published as a conference paper at ICLR 2024
Proof. By Propositions 12 and 15 with parameters n, k K, q p. and define
1/vP. if
N
N + CNV log
Nmin
Ci N
N Mmxx
(A.178)
then we have
PIll An
> s]
(A.179)
Naax
C2
AI A
I|| >
K₁ = n]]p[K
n| +
(A.180)
N
n= Niin
11))
Nmax
exp
1 CCCVND
8
PlKK n)
(A.181)
N2
1 1
N
Cxpxp
if C..
In order to cancel the most terms, we choose
Nmax p
=
(A.182)
In order to assure ourselves that this choice still has S C7 (so that we can use the first case for all n2,, we have
Nmax
==
(A.183)
P
N/K + CIV. log
20
(A.184)
N/K + I
20,
(A.185)
3
N/K
2
clos 2
(A.186)
N
voclos
V
(A.187)
d
N when VGCs
S
1
(A.188)
d
Of course, this condition is assured by Assumption 2. Now that we have this, we know falls in the first, and so we have
N
A Av
1|| CIV d
(A.189)
-n Q 1/))
Nmax
Cxexp
PlKK n).
(A.190)
N2
C 1 iss0
Nui
Csexp
C2 N2
Nmax
OCRCssNmaake
Czex
P(Ka n]
(A.191)
CROss C25svo/D
N Nain
C2
N
waax
-n (
Nuax
1)
C exep
2
2
PP[K n)
(A.192)
N2
n==Nmin
C2
< Ca exp
(A.193)
2
K
N2
where the last inequality follows from the exact same argument in Proposition 10.
33

---

Published as a conference paper at ICLR 2024
A.2.3 GENERIC CONCENTRATION INEQUALITIES
In this subsection we prove the base-level concentration inequalities used throughout the proofs in this paper.
Binomial concentration.
Proposition 12. There exist universal constants C1, C2 O such that the following holds. Let n, E L. For each E [k), let B, Bin(n, 1/k), such that the B₁ are identically (marginally) distributed but not necessarily independent binomial random variables. Let E. be event. Then for any E [k), we have
In/k+Ciynlogy Σ
C2
P S
F|E| B; - B
b] + n2
(A.194)
Proof. We have
P E| B,]]
P B, = bbP[B: =b].
(A.195)
b=0
Each B, is unconditionally distributed as Bin(n, 1/k). By union bound and Hoeffding's inequality (Vershynin, 2018, Theorem 2.2.6), we have
2t
P B
n/k 2 t S 2 exp
(A.196)
Inverting this inequality obtains that there exists some (simple) universal constants C₁, C2 > 0 such that
C2
P [ n//k 2 C1 CLvno log S n
(A.197)
Thus, if we define
6min
Ci V n logn
bmax
=
CCV log
(A.198)
k
then we have
P
P[E| B; = b]P[B: = bl
(A.199)
b=
harin
P B, = P b+ +
P
B, = b
P(B: b
(A.200)
bm0
S1
SCo/n3
h +mncett
<1
SC2/n²
hmna
+
P[E| B; b]pp[: = b]
(A.201)
bran
mnin
C2
hans
P[E
B: blP[B. bj
(A.202)
3
b=0
b=
b=mmax +1
b=
hans
<
P[E B = b]P[B: = b]
(A.203)
7
b=0
C
hanax
P[E
B; = BJP[BB = bj
(A.204)
112
3
b=
2C²
Σ
b).
(A.205)
n
b=bmi
34

---

Published as a conference paper at ICLR 2024
Remark 13. Notice that simple adaptation of this argument turn the additive probability Cs/n? into C3/n" for any positive integer 2 E N (where C32 depends on 2). However, trying to replace with CGe C'n is more difficult.
Gaussian concentration.
Proposition 14. There universal constants C1, C22 C3 > 0 such that the following holds. Let n, E N, and let M E Rox be such that Mij wiid No,, x2) Then
8/0 V9
M| > s] S C1 exp
-n
C2V1
1/
Vs > x + +V
(A.206)
-q {a(
-1111 11))
10/
8//0 vn
[[[ M > 8) S C1 exp
A&> x (Vn Cav
(A.207)
Csv4
Proof. Define M M, so that My midd (00,1) 1). By Vershynin (2018, Example 2.5.8, Lemma
3.4.2), we see that each row M, has Orlicz norm M, S C₁ for some universal constant C1 > 0.
By Vershynin (2018, Theorem 4.6.1) we have for some other universal constant C2 > that for all
t > 0,
0 - CC2CV + S mi (M) S 01 (M) I Vq CC22(V++
(A.208)
with probability at least - Defining C3 = C2C₂ and noting that ||-| = 01 (1)), we have with the same probability that
M v - C3 (Vn+t))
(A.209)
Simplifying, we obtain
M v (Vū++
(A.210)
M v4s (Vi++t)
(A.211)
M - ~ s Ca (Vi+t)
(A.212)
|M] s x (V7 Co (vn+t))
(A.213)
Define by
8/I /
(v7 Cs (Vm+t)) =
Vn.
(A.214)
C2
Note that the range of validity is
+> 3> s> + Covm).
(A.215)
For in this range, have
( (-1//-- vi
vn (n)))) 2exp -n
s/2-vv v
- C3v -11 1 17)) (A.216)
M| || > s) < 2 exp
C3
The other inequality follows from applying this inequality to M..
Proposition 15. There are universal constants C1 C2 C3 > 0 such that the following holds. Let n, E N, and let M E Roxn be such that Mij miid. (0, Then
P M M
qr
s]
(A.217)
C exp
PBCosnga - -1
if Ca S Cz
(A.218)
91 exp
-n
CsCayas V8 15-11
ifs 2 Czaz2.
35

---

Published as a conference paper at ICLR 2024
Proof. Define M = _ M,, so that My wi.d Mo0, 1). By Vershynin (2018, Example 2.5.8, Lemma 3.4.2), see that each row has Orlicz norm | S C₁ for some universal constant C₁ > 0.
By Vershynin (2018, Eq. 4.22) we have for some other universal constant C2 > 0 that for all t >> 0,
- M M I
- Cf max 06. where C= /n +
(A.219)
Vq
with probability at least
2e 2e- Simplifying, we obtain
1 M M- - O9 max (8, 8²}
(A.220)
| M- - q1 || - 09 q max {6, 8²)
(A.221)
|| M) T A-MM - q1 - 07q max { (8 82)
(A.222)
||> 2 M T M - -q - maq f 6, 8²)
(A.223)
M M - qx- 2 ||| s G9Tz² max (0, 8²).
(A.224)
Now from simple algebra and the fact that 1, we have
max {0,8*} = = 02 I VI Vn
(A.225)
max {8,02) = 8² = 12 C5 VI Vn.
(A.226)
Now define s 2 0 by
Ciqx" max (6, 8²}.
(A.227)
Thus in the first case we have
= C:C (vi++) = t= OPC2/002 Vn,
(A.228)
and in particular the first case holds when
max (6, = 5 = Meanwhile, in the second case, have
Vn= = ChCovnqa S S Cggo2.
(A.229)
=
CFC3(vn = C1C2x
Vn
(A.230)
where we obtain only one solution to the quadratic equation by requiring / 2 0, and in particular the second case holds when
max max{o,0²}= = + t CZ Vq 49 Vn + 3 C2qx2.
(A.231)
Thus we have
P|[M M- aM--z221 > s]
qr²1|
(A.232)
2exp
CCCa V1
CCGCopqqu' < S Czqx²
(A.233)
2)
2exp
CiC2z V8 v
1
if 2 C2
2exp
n
CC22nggz 3
if C /パ < < CZ
(A.234)
2exp
-n
C.Coynn V3-14 1]
2
if 2 C99z2.
36

---

Published as a conference paper at ICLR 2024
A.3 COMPANION TO SECTION 2.4
In this section, we justify the scaling applied to VR in Section 2.4 and supply the discretization scheme.
First, suppose that Z satisfies (2.1), and Z: = z + o W, where W is a standard Gaussian matrix, so that Z satisfies (2.2) with noise level og > O. Let qe be the density Z, Theoretical analysis from (Lu et al., 2023) and empirical analysis from (Song & Ermon, 2019) demonstrates that under generic conditions, we have that
I (Z:)l22 1 o.
(A.235)
ignoring all terms in the right-hand side except for those involving os. On the other hand, from the proof Theorem 3, we obtain that VV--VRRM (Z))(,) Zz) has constant (in o₂) magnitude with high probability. Thus, in order to have them be the same magnitude, we need to divide - VR (z) by of to have it be a drop-in replacement for the score function, as alluded in Sections 2.3 and 2.4.
R
Second, we wish to explicitly state our discretization scheme given in Section 2.4. To wit, we provide a discretization scheme that turns the structured diffusion ODE (2.15) into its discretized analogue; the other discretization of the structured denoising ODE (2.11) occurs similarly. To begin with, define the shorthand notation
f(t. (()) = VHE ( t)),
(A.236)
so that we have
( = ftt, (())) dt.
(2.11)
2t
Fix I., and let 0 < <4 < 12 <... < 4 = T, such that 4 is small. (These will be specified shortly in order supply the discretization scheme.) A suitable first-order discretization is given by
1411 - t f(te, Y
y ~ Y
(A.237)
21e
Thus it remains set t t,,..., tu such that
te+1 te
K
(A.238)
21
for some constant we observe that we must set
t++1 = (1 2x)te.
(A.239)
so that the time grows exponentially the index. The reverse process time decays exponentially in the index, which matches practical discretization schemes for ordinary diffusion models (Song & Ermon, 2019). Finally, we have T ty = (1 + 2x)(tt (t. so that t [I42k)
37

---

Published as a conference paper at ICLR 2024
B ADDITIONAL EXPERIMENT DETAILS
B.1 EXPERIMENT DETAILS AND CLARIFICATIONS
In all setups, as is standard practice (He et al., 2022), we append a trainable class token z E Rd after masking and linear projection, namely,
fore X) = (z W X + E
(B.1)
Everything else goes through with Z having N + instead of N tokens, indexing the N image tokens' intermediate representations as zf, etc., and indexing the class token intermediate represen- tation as dls At the end, the post-processing map is
grost (Y W L 1:N
(B.2)
where Y tokens, i.e. the second through last columns of Y Thus, the class token z (and its representation Rels, i.e. the first column of Z) are neither masked or reconstructed.
I:N E MRdxN is the matrix whose columns are the columns of Y corresponding to image
For training using masked autoencoding, we follow a modified recipe of (He et al., 2022). We mask a fixed percentage y E [0, 1] of randomly selected tokens X; that is, we randomly set (1 - n) image tokens æ; 0 DR. obtaining X. Then X the input to the encoder. Unlike in (He et al., 2022), the decoder receives special treatment, and operates on all token representations. The loss computed only on the masked image patches. Refer to He et al. (2022) for more MAE implementation details.
For fine-tuning using supervised classification, we use the representation zels of the vestigall class token the feature used for classification. Namely, obtain the unnormalized log- probabilities for the classes as u Whead LN(z.ts), where LN a trainable layer-norm and Whead E Cxdd is trainable weight matrix, where C' is the number of classes. The output u E RC the input to the softmax cross-entropy loss. All model parameters are trainable during fine-tuning, while in linear probing only the weight matrix Whead is trainable (and in fact learned via full-batch logistic regression).
In all training setups, average the loss over all samples in the batch.
We pre-train CRATE-MAE on ImageNet- (Deng et al., 2009). We employ the A optimizer (Loshchilov & Hutter, 2019). We configure the learning rate x 10-3, weight decay as 0.1, and batch size as 4, 096.
We fine-tune and linear probe our pre-trained CRATE-MAE on the following target datasets: CI- FARIO/CIFARI00 (Krizhevsky et al.. 2009), Oxford Flowers- (Nilsback & Zisserman, 2008), Oxford-IIIT-Pets (Parkhi et al., 2012). For each fine-tuning task, we employ the AdamW opti- mizer (Loshchilov & Hutter, 2019). We configure the learning rate as 5 x 10-5, weight decay as 0.01, and batch size as 256. For each linear probing task, we use the linear probing function- ality in Scikit-Learn (Pedregosa et al., 2011). 1). For each evaluation we choose several regularizers C' E {109, 10a, 102, 100, 10a, 10 train a logistic regression model on features from the whole dataset, and choose the logistic regression model with the best performance. All numbers are ported on the test sets.
For experiments,
the model configurations reported in Table 1. From this table, there two unspecified hyperparameters, namely À and 17. In all experiments we fix 1 = 0.1, sincce only multiplies with trainable matrix and X. In numerical experiments figures are generated from models with hyperparameter À = 5.0, though the difference numerical performance and figure quality is marginal between the two settings (Table 6).
À 0.5, while
To allow transfer learning, in all training and evaluations setups we first resize our input data to 224 height and width. For data augmentations during pre-training and fine-tuning, we also adopt several standard techniques: random cropping, random horizontal flipping, and random augmentation (with number of transformations
and magnitude of transformations m 14).
B.2 PYTORCH- PSEUDOCODE
38

---

Published as a conference paper at ICLR 2024
Listing 1: PyTorch-Like Code for MSSA and ISTA
class ISTA:
initialization
def _init_ (self, dim, hidden_dim, dropout O., st
Lambd=0.1):
super(). iniitt ()
self.weight. Paraaaatte (torch Tensor (dim, dim)
with torch corch.no_graa()
hntt.kaiming_unifrrmm. (sel weight)
setf.step_size step_size
self.lambd = lambd
forward pass def (sel,,,
x1 F.linear (x, welf.weight, bias-None)
grad_ = F. E.lincar (x1, self. wolf.weight.t(). bias-None)
grad_2 = Filinar (x, self. welf.weight.t() bias-Nenc)
grad_update step.etep_size (grad_2 grad_l) self step_stepssize
self.lambd
output - (ereuux + grad_update)
return output.
class MSSA:
initialization def _ii___nit_____________ (self, dim, heads 8, dim_head 64, dropout 0.):
inner_dim - dim_head heads
project_out not (heads == and dim_head dim)
self.heads heads
self.scale dim head ** -0.5
self.attend Softmax (dim = -1) sclf.dropout Dropout (dropcut) self.qkv Linear (dim, inner_dim: bias=False)
self.to_out Sequential (Linear (inner_ding dim), Drepout (dropout) if project_out else mn. Indetiiity)))
forward pass def fowar, x):
rearrange
(x), "b d)
d', self.
acts matmul (w, w.transpose -21) self.scale
attn atlend
attn self. self.dropout (attn)
out matmul (attn, out rearrange (out, (b -> bn (h d)
return self.to_ous (out)
Listing PyTorch-Like Code for CRATE-MAE Encoder
class :AATE__nccder:::
initialization
def innii_ttstl__ (self, dim, depth, heads, dim_head, mip_din, dropout
0.):
sei fayers []
self.depth depth
for in range (depth)
Lelf.layers.exx ([LayerNorm (dim), MSSA (dim, heads, dim head, dropout]
sels. extend (LayerNorm (dim), ISTA (dim, mlp dim,
dropout)
forward pass
def forward x):
for Inl, attn, 1n2, if in sclf. layers: x_ attn (1n1 (x))
(1n2
return x
39

---

Published as a conference paper at ICLR 2024
Listing 3: PyTorch-Like Code for CRATE-MAE Decoder
class CRATE_ Decoder
initialization
def _init_____ _init__self, dim, depth, heads, dim_ mip_dim, dropout
0.): # define layers
selt.
sel self.depth - depth
for in range (depth):
Self.layers.com out out_ fes bias=False)])
[layerNorm (dim), Linear (in fin_ffaature-ddn,
lels. ([LayerNerm (dim), MSSA (dim, heads, dim head, dropout)
forward pass
def (sewf, x):
for In1, f_linear, 1n2, attn in self. layers: x f_linear (1n1 (x))
1n2 (Itn (x_))
return x
B.3 VISUALIZATION METHODOLOGY
In this subsection we formally describe the procedures we used to generate the visualizations used evaluate the segmentation property of CRATE-MAE in Figure 7. Much of this evaluation the same as in Yu et al. (2023b), which initially demonstrates the emergent segmentation properties of white-box architectures.
B.3.1 PCA VISUALIZATIONS
We recapitulate the method to visualize the patch representations in Figure 7a using PCA from Amir et al. (2021); Oquab et al. (2023); Yu et al. (2023b).
We first select J images that belong to the same class, {X, 1 and extract the token representa- tions for each image at layer f, i.e., [=f]....., j.N for j [J]. In particular, 21,, represents the 7th token representation at the fth layer for the jth image. We then compute the first principal com- ponents of 2' representation for the i-th token of Xj. i.e., $ = [(U)
,
24 ........ {, and use Σ to denote the aggregated token
(5, TIT e R(r-K)11,
We denote the first eigenvector of the matrix
7 2 by uo and compute the projection values as x. x
{ax 3e where (x) = 0. < is the hard-thresholddigg function. We then se-
lect subset of token representations from Z with ax((u,, 2f,1)) > O. which correspond to non-zero
projection values after thresholding, and we denote this subset 2. S 2. This selection step is used to remove the background (Oquab et al., 2023). We then compute the first three right singular vectors of Z, with the first three eigenvectors of the matrix Z:Z. denoted as {u, u2 uz). We define the RGB tuple for each token as:
[rjáá 9i.s, bjal] = [(u,, L. (u2, 2),) (u3, 2,,.
i E [N], j [M, 11, E Z..
(B.3)
Next, for each image X; compute R, Gj, Bj. where R, frg,,...,, for Gj and Bj). Then reshape the three matrices into VN x VN and visualize the "principal components" of image X, via the RGB image (Ry, Gj, Bj) E RR3/NNVV
Rdx1 (similar
B.3.2 VISUALIZING ATTENTION MAPS
We recapitulate the method to visualize attention maps in Abnar Zuidema (2020); Caron et al. (2021).
40

---

Published as a conference paper at ICLR 2024
Table 3: Top-1 classification accuracy of CRATE-MAE- when pre-trained on ImageNet-1K and fine- tuned on classification for various datasets. We compare a fine-tuned model which was pre-trained on the MAE task with a model trained from scratch on classification ("random inif") using exactly the same experi- mental conditions. Our results show that the representation learning occurring during pre-training substantially improves performance downstream tasks.
Classification Performance
CIFAR 10
CIFAR 100
Oxford Flowers
Oxford-Pets
CRATE-MAE-Bars (trained)
96.8
80.3
78.5
76.7
CRATE-MAE-Base (random init)
85.1
58.8
38.0
28.8
Table 4: Average reconstruction loss over the training and validation sets of ImageNet-1K 1K for both CRATE-MAE- MAE-Base and VIT-MAE-Base We see that the performance of CRATE-MAE-Base while bit worse than ViT-MAE-Base, obtains promising performance the challenging masked autoencoding task.
Reconstruction Loss
Training Loss
Validation Loss
CRATE-MAF-Base ViT-MAF--aas
0.265
0.302
0.240
0.267
Table 5: Top-1 classification accuracy of CRATE-MAE --aee when pre-trained on ImageNet-1K and lincar probed for classification CIFAR-10, when pre-trained using different mask percentage (i.e., number of masked tokens each sample). This shows that CRATE-MAN MAE models with 75% of tokens masked during training tend to have the most structured representations that are useful for downstream tasks, an empirical conclusion that echoes He et al. (2022), but a wide range of mask percentages result in good representations. (Note: This table uses the hyperparameter setting À = 5.0.)
Classification Accuracy
25% Masked
50% Masked
75% Masked
90% Masked
CLATE-MABB M-A-BASe Base
69.78
75.97
75.99
73.45
For the 1tth head at the gth layer of the encoder of CRATE-MAE, compute the self-attention map A E N defined as follows:
A 1
exp((UE+ 2f, 2f ((a))) A6, Σ5 (U 24/, =()))
E CRN.
where
(B.4)
A
where zf is the gth layer representation of the class token.
For each image, we reshape the attention matrix A for the penultimate layer I, into /N x VN matrix and visualize the heatmaps as shown in Figure 7b. For example, the 7th row and the jth column element of each heatmap in Figure 7b corresponds to the mh component of A where m (i - 1) /N + j. In Figure 7b, for each image we select one attention head & of CRATE-MAE and visualize the attention matrix A/-1
B.4 ADDITIONAL EXPERIMENTS
In this section we perform more experiments to explore properties of the CRATE-MAE architecture.
First, in Table 3. we compare the fine-tum tne-tuning performance of an CRATE-MAE MAE-Bass model with MAE pretraining against an PRATE-MAF-BaSS model with no pretraining at all (i.e., randomly initialized). We apply the same fine-tuning process to both models, and we observe a massive disparity in per- formance, where the pre-trained model succeeds while the randomly initialized model performs poorly. This indicates that the organized representations of a pre-trained CRATE-MAN MEE---MEBBaSS model a strong starting point when fine-tuning for downstream tasks.
Next, in Table 4, we evaluate the reconstruction loss (measured in mean-square error) of CRATE- MAE-Base versus VIT-MAE-Basee evaluated on the training and test sets of ImageNet-IK. We ob- serve that while CRATE-MAL MaE-Base performs slightly worse at masked reconstruction, the perfor- mance is still very reasonable.
41

---

Published as a conference paper at ICLR 2024
Table 6: Top-1 classification accuracy of CRATE-MAE-Bas when pre-trained ImageNet-Ik 1K and linear probed for classification various datasets, when pre-trained using different This shows that CRATE- MAE models perform reasonably well at CRATE-MAF- scale across different values of
Classification Accuracy
0.5
TRATE-MAE-Bas.
83.33
80.87
75.99
Masked
VIT-MAE
CRATE-MAE
Original
Masked
VIT-MAE
CRATE-MAE
Original
Figure 8: More instances of parity between CRATE-MAE- toencoding task. Echoing the message of Figure 6, we find that CRATE-MAL Base and ViT-MAE Base have similar performance on the masked autoencoding task, even as the CRATE-MAF-BasS model is significantly more parameter-efficient.
ViT-MAE- Base in the masked au-
In Table 5, check the downstream performance and feature learning of CRATE-MAE-Base models when varying the mask size. We use test accuracy of linear probing CIFARI0 as a proxy for the quality of the learned features. Our results show that the performance is maximized when the number of masked tokens is 75% of the number of total tokens, i.e., when 75% of tokens are maxed out. This confirms the experiments in He al. (2022).
In Table 6, we check the downstream performance and feature learning of PRATE-MAE-BaR models when varying the hyperparameter We again use test accuracy linear probing on CIFARIO 10 as proxy for the quality of the learned features. Our results show that the performance is maximized when À = 0.1, but models perform reasonably well at CRATE-MAR -Bss scale.
In Figure 8 we provide more examples of the masked autoencoding efficacy of both CRATE-MAE- Base and ViT-MAE-Base. Our results confirm those of Figure 6, namely CRATE-MAE: SAATE-MA-BBaS -aae achieves parity with the much larger ViT-MAE-BBan
the challenging masked autoencoding task.
In Figure we provide more examples of the linearity of the representations within CRATE-MAR MAE- Base. We see that correlate strongly with semantically meaningful patches of the input image. These results extend Figure 7a.
wide variety of images, the first three principal components of each class
42

---

Published as a conference paper at ICLR 2024
Figure 9: More examples of linearized representations CRATE-MAE- Base. Echoing the message of Figure 7a, we find that CRATE-MAE -RAA--MAE-Base has a linear feature space. In particular, the first three principal components strongly correlate with the main semantic content of the image.
More examples of interpretable attention maps in CRATE-MAE-Bane Echoing the message of Figure 7b, we find CRATE-MAE-BEE has numan-interpretable attention maps which semantically segment the foreground input images.
Figure 10:
In Figure 10, we provide more examples of the interpretability of selected attention outputs within SRATE-MAEE-asSE We see that over wide variety of images, the attention map succeeds at captur- ing many semanties of the input image. These results extend Figure 7b.
In Figure 11 we apply (nearly) the same methodology involved in constructing Figure 7, qualitative demonstration of the feature quality of the CRATE-MAF Base encoder, to evaluate the feature quality of the ViT-MAE-Base encoder. One necessary difference is that to evaluate the attention maps, the VIT does not have the UiK] matrices that CRATE does, but instead has three sets of matrices Q(K)' Kt and V [k] thus, construct the attention maps via the following equation: [K].
exp(( 2 ++=()))
Ab. = EN (oxp((((e) Q =1)
(B.5)
43

---

Published as a conference paper at ICLR 2024
(a) Visualizing PCA of ViT-MAE-Bar Base.
(b) Visualizing attention maps of ViT-MAE-Bas
Figure 11: comparison CRATE to ViT-MAE in the setting Figure 7. We use the visualizations of PCA on the token representations and attention maps, introduced in Figure 7, to qualitatively evaluate the representation quality of the ViT-MAE-Base Base. By comparing this figure (Figure 11) and Figure 7, observe that CRATE attention maps contain more clear semantics than those from VřT-MAE-Base, while both CRATE-MAF- and ViT-MAE-Base have nearly-linear representation spaces wherein semantic concepts correspond to the first three principal components.
Overall, Figure 7 and Figure 11 demonstrate that, at least empirically, the attention semantics in CRATE-MAE-Base are significantly better and clearer than ViT-MAB-Base The reason that CRATE- MAE models have semantically meaningful attention maps may be due to white-box design, namely the fact that CRATE-MAE we have iik] = Kix k] Vik = Uikj indeed, the fact that
Vik Ukk yields semantically meaningful attention maps has been shown in other work, albeit different setting (Yu et al., 2023b). The reason that the ViT-MAE has semantically meaningful attention maps, albeit not as clear and worse than CRATE-MAE, may be due to several different factors such as using the class token as a register (Darcet et al., 2023). Nevertheless, the features in both models have linear structures, or at least each class' representations have three principal components which correlate strongly with the semantics of class.
setting Oxx]= K
44


Reviews:
[{"review_id": "zwkNJKCTti", "reviewer": "Reviewer_pN5Y", "summary": "This paper generalizes the white-box design of transformer, i.e., CRATE, to the unsupervised representation learning context. The author finds out that the gradient on the rate distortions term $R^c(Z | U_[K])$ plays a similar role as the gradient for the score function with the noised input $\\tilde Z$, which points towards the closest point to $\\tilde Z$ on the data distribution support. Thus they construct a masked auto-encoder using CRATE backbone and achieves fine results on the representation learning tasks.\n\n===============\nThe author address most of my concerns and I will increase my score", "questions": "1. Which dataset does Fig.4 belongs to? Does the model have the similar patterns as Fig.4 on other datasets evaluated in this paper?\n2. Empirically, people think that the attention map $Q/K$ plays a different role as the mapping matrix $V$ and they'd better not be set to be the same. However, in this papers' theoretical framework, they can be set to the same parameter $U$. If change the mapping matrix from $U$ to $V\\neq U$, will the performance of CRATE-MAE change a lot? If no, what's the main reason why CRATE-MAE has such property?\n3. Will the layer normalization influence the rate deduction process, or it's just for making the training process more stable or other reasons?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. This work generalizes the white-box transformer model to the unsupervised representation learning task, which is a novel attempt to both the theory and empirical community.\n2. The visualization results are quite impressive and show that CRATE-MAE can reconstruct the original data well, and Fig. 4 roughly shows that the compression measure and sparsity measure match the theory setting.", "weaknesses": "1. I think the comparisons between CRATE-MAE-Base amd MAE-Base are not fair. I understand that the empirical evaluations are not optimally engineered and actually the visualization results are every good. However, I still think that the author should compare CRATE-MAE and MAE with (almost) the same amount of parameters and different performance, or alternatively, (almost) the same performance and different amount of parameters, and then to compare these two models.\n2. What's the choice of LASSO coefficient hyperparameter $\\lambda$, and the step size of discretization $\\kappa, \\eta$? Are they chose carefully and is the model sensitive to them?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fC6uNTQyzB", "reviewer": "Reviewer_vMuz", "summary": "The paper extends the white-box models from supervised learning to unsupervised (or self-supervised) representation learning. And in particular, it trains a Masked Autoencoder (MAE) to learn transferrable representations to downstream classification tasks. Besides, it shows interesting visualizations to demonstrate learned representations are of emerging semantic properties.", "questions": "* The paper lacks a fair comparison with MAE, especially on downstream tasks. Table 1 lists that MAE-base has more parameters and it could partially explain why Crate-MAE has a higher reconstruction loss. So what would a fair comparison (in terms of model parameters) look like? I think it is very easy to train MAE with a smaller encoder/decoder pair given the open-sourced code.\n* How are the evaluations done in Table 2? Are they using the encoder only? Are they with fully-visible inputs?\n* Table 4 again lacks a comparison of a similarly-sized MAE (how it performs on the same dataset as the mask ratio changes). It is not clear a conclusion from ImageNet classification can be transferred to CIFAR.\n* Figure 4: how does it compare with a supervised encoder/decoder trained on ImageNet? I want to know the compression and sparsity behavior is a result of unsupervised learning, or a result of the architecture design. The same applies to Figure 6 and 7.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ Learning unsupervised representations for white-box models is a natural topic to study after white-box on supervised models, and could be of interest to members of the community.\n+ The paper did a good job in presentation, making both the approach and the experiments easy to follow.\n+ I like the visualizations in the end, which is a more intuitive addition to the numerical comparisons in the table.", "weaknesses": "- I see a lot of similarities to the main white-box paper (White-Box Transformers via Sparse Rate Reduction) which is already out there for supervised learning. I want more justifications for the meaningfulness of the current work apart from doing unsupervised learning. \n- I am not convinced what's the advantage of models being white-box here, especially whether it has synergies with unsupervised learning. The paper explains a lot about what's done, but I don't see a strong motivation of why it is done (especially since the introduction is less of a story but more of a break-down of context and contributions).\n- While the explorations are interesting, I don't think the claims are backed up well by experiments. This is my biggest concern and would like to raise them by asking questions. So please see below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1Ck4xXZxno", "reviewer": "Reviewer_86oH", "summary": "This paper proposes a white-box diffusion model and unifies it with data compression under a single framework. Based on this, the method proposed by the author has achieved results comparable to the state-of-the-art, which validates their proposed theory.", "questions": "Please refer to my comments in the weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Novelty in Approach: This paper is the first attempt to turn the diffusion model into a white-box network, and it has achieved good results. \n\nVersatility: Its methodology and conclusions can be used as good reference for subsequent research on white-box neural networks.", "weaknesses": "1. In Sec 2.2, the authors intend to learn representations Z, and hope that the results learned are low-dimensional, sparse, bijective, etc. \nIf the method proposed by the authors is a white box, then these properties of Z should be verifiable through experiments. \nTherefore, the authors should provide experimental results of representations Z to support their theory. \n\n2. I checked the provided Pytorch code and find that MSSA and ISTA are composed of Linear layer, which implies large GPU Memory consumption.\nSo I'm wondering whether this method can be extended to large images, just like stable diffusion. \nIn addition, can the authors provide results from larger datasets? The images from CIFAR and ImageNet-1k are too small. \n\n3. How effective is this network at unconstrained image generation? \nIn other words, if the learning target is pure noise, which may not be viewed as an image compression task, would this method still work?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qr20GIAU2o", "reviewer": "Reviewer_2DR1", "summary": "This paper introduces CRATE-MAE, a novel white-box deep network architecture designed for large-scale unsupervised representation learning. Unlike traditional networks, CRATE-MAE is rooted in the mathematical connection between diffusion, compression, and masked completion. Each layer of this architecture has a clear, interpretable role, transforming data into structured representations and vice versa. The study's key contribution is adapting the white-box design for unsupervised learning, a notable departure from its typical supervised applications. Empirically, CRATE-MAE outperforms traditional models on large imagery datasets with 30% fewer parameters, while offering structured and semantically meaningful representations.", "questions": "I have no questions about the methodology part. Just please add more quantitative results to paper.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. **Theoretical Depth and Scientific Rigor**: The research stands out for its robust theoretical foundation, seamlessly intertwining denoising-diffusion models and information theory with white-box models. The model's design is intricately tied to its theoretical underpinnings, exemplifying the paper's scientific precision and thoroughness.\n\n2. **Problem Significance**: By addressing representation learning in high-dimensional data and delving into the untapped potential of white-box models in unsupervised settings, the paper carves a significant niche in the contemporary machine learning domain. Indeed, it paves the way for new avenues of exploration for the broader ML community.", "weaknesses": "Firstly, I'd like to clarify that my emphasis is not solely on state-of-the-art results. My questions regarding the experiments stem from the belief that robust ideas and arguments deserve to be bolstered by thorough experiments.\n\n1. **Evaluation Concerns**: The introduction of the CRATE-MAE architecture in the paper falls short in offering a comprehensive quantitative analysis when compared to established benchmarks like MAE or contrastive methods. The results presented in Table 2 seem somewhat restrictive, making it challenging for readers to gauge the model's efficacy in relation to others.\n\n2. **Local Self-Supervised Learning Comparison**: From a broader perspective on self-supervised learning (SSL), this paper could be classified under local SSL, emphasizing layer-specific objectives and training. Although this area might be less traversed, incorporating findings from related works, such as [A], would enhance the paper's credibility.\n[A] Siddiqui, Shoaib Ahmed, et al. \"Blockwise self-supervised learning at scale.\" arXiv preprint arXiv:2302.01647 (2023).\n\n3. **Absence of Linear Probing Results**: Omitting linear probing results restricts the paper from showcasing the practicality and caliber of the representations derived using CRATE-MAE.\n\n4. **Dataset Limitations**: The study's dependence on a confined dataset for classification raises concerns about the breadth of its applicability and potential generalization to diverse scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)