SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yaodong Yu<sup>1</sup> Sam Buchanan<sup>2</sup> Druv Pai<sup>1</sup> Tianzhe Chu<sup>1</sup> Ziyang Wu<sup>1</sup> Shengbang Tong<sup>1</sup> Benjamin D. Haeffele<sup>3</sup> Yi Ma<sup>1</sup>
<sup>1</sup>University of California, Berkeley <sup>2</sup>TTIC <sup>3</sup>Johns Hopkins University
In this paper, we contend that the objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a mixture of low-dimensional Gaussian distributions supported on incoherent subspaces. The quality of the final representation can be measured by a unified objective function called sparse rate reduction. From this perspective, popular deep networks such as transformers can be naturally viewed as realizing iterative schemes to optimize this objective incrementally. Particularly, we show that the standard transformer block can be derived from alternating optimization on complementary parts of this objective: the multi-head self-attention operator can be viewed as a gradient descent step to compress the token sets by minimizing their lossy coding rate, and the subsequent multi-layer perceptron can be viewed as attempting to sparsify the representation of the tokens. This leads to a family of white-box transformer-like deep network architectures which are mathematically fully interpretable. Despite their simplicity, experiments show that these networks indeed learn to optimize the designed objective: they compress and sparsify representations of large-scale real-world vision datasets such as ImageNet, and achieve performance very close to thoroughly engineered transformers such as ViT. Code is at https://github.com/Ma-Lab-Berkeley/CRATE.
In recent years, deep learning has seen tremendous empirical success in processing massive amounts of high-dimensional and multi-modal data. Much of this success is owed to effective learning of the data distribution and then transforming the distribution to a parsimonious, i.e. structured and compact, representation [39, 50, 52, 62], which facilitates many downstream tasks (e.g., in vision, classification [23, 40], recognition and segmentation [25, 38, 77], and generation [31, 65, 66]). To this end, many models and methods have been proposed and practiced, each with its own strengths and limitations. Here, we give several popular methods a brief accounting as context for a complete understanding and unification that we seek in this work.
Transformer models and self-attention. Transformers [28] are one of the latest popular models for learning a representation for high-dimensional structured data, such as text [28, 30, 37], images [40, 75], and other types of signals [48, 57]. After the first block, which converts each data point (such as a text corpus or image) into a set or sequence of tokens, further processing is performed on the token sets, in a medium-agnostic manner [28, 40]. A cornerstone of the transformer model is the so-called self-attention layer, which exploits the statistical correlations among the sequence of tokens to refine the token representation. Transformers have been highly successful in learning compact representations that perform well on many downstream tasks. Yet the transformer network
1{yyu,yima}@eecs.berkeley.edu, {druvpai,chutzh,zywu,tsb}@berkeley.edu
sam@ttic.edu
3 bhaeffele@jhu.edu
Figure 1: The 'main loop' of the CRATE white-box deep network design. After encoding input data  $X$  as a sequence of tokens  $Z^0$ , CRATE constructs a deep network that transforms the data to a canonical configuration of low-dimensional subspaces by successive compression against a local model for the distribution, generating  $Z^{\ell +1 / 2}$ , and sparsification against a global dictionary, generating  $Z^{\ell +1}$ . Repeatedly stacking these blocks and training the model parameters via backpropagation yields a powerful and interpretable representation of the data.
architecture is empirically designed and lacks a rigorous mathematical interpretation. In fact, the output of the attention layer itself has several competing interpretations [68, 78]. As a result, the statistical and geometric relationship between the data distribution and the final representation learned by a transformer largely remains a mysterious black box.
Diffusion models and denoising. Diffusion models [22, 34, 41, 43, 44] have recently become a popular method for learning the data distribution, particularly for generative tasks and natural image data which are highly structured but notoriously difficult to effectively model [3, 5]. The core concept of diffusion models is to start with features sampled from a Gaussian noise distribution (or some other standard template) and iteratively denoise and deform the feature distribution until it converges to the original data distribution. This process is computationally intractable if modeled in just one step [61], so it is typically broken into multiple incremental steps. The key to each step is the so-called score function, or equivalently [13] an estimate for the "optimal denoising function"; in practice this function is modeled using a generic black-box deep network. Diffusion models have shown effectiveness at learning and sampling from the data distribution [56, 60, 65]. However, despite some recent efforts [81], they generally do not establish any clear correspondence between the initial features and data samples. Hence, diffusion models themselves do not offer a parsimonious or interpretable representation of the data distribution.
Structure-seeking models and rate reduction. In both of the previous two methods, the representations were constructed implicitly as a byproduct of solving a downstream task (e.g., classification or generation/sampling) using deep networks. However, one can also explicitly learn a representation of the data distribution as a task in and of itself; this is most commonly done by trying to identify and represent low-dimensional structures in the input data. Classical examples of this paradigm include model-based approaches such as sparse coding [2, 29] and dictionary learning [17, 21, 47], out of which grew early attempts at designing and interpreting deep network architectures [18, 32]. More recent approaches build instead from a model-free perspective, where one learns a representation through a sufficiently-informative pretext task (such as compressing similar and separating dissimilar data in contrastive learning [45, 69, 80], or maximizing the information gain in the class of maximal coding rate reduction methods [6, 46, 55]). Compared to black-box deep learning approaches, both model-based and model-free representation learning schemes have the advantage of being more interpretable: they allow users to explicitly design desired properties of the learned representation [46, 55, 63]. Furthermore, they allow users to construct new white-box forward-constructed deep network architectures [11, 55, 59] by unrolling the optimization strategy for the representation learning objective, such that each layer of the constructed network implements an iteration of the optimization algorithm [11, 53, 55]. Several recent works [71, 74, 76] consider the connections between transformer architectures [28] and unrolled optimization. Unfortunately, in this paradigm, if the desired properties are narrowly defined, it may be difficult to achieve good practical performance on large real-world datasets.
Our contributions, and outline of this work. In this work, we aim to remedy the limitations of these existing methods with a more unified framework for designing transformer-like network architectures that leads to both mathematical interpretability and good practical performance. To this end, we propose to learn a sequence of incremental mappings to obtain a most compressed and sparse representation for the input data (or their token sets) that optimizes a unified objective function known as the sparse rate reduction, specified later in (1). The goal of the mapping is illustrated in Figure 1. Within this framework, we unify the above three seemingly disparate approaches and show that transformer-like deep network layers can be naturally derived from unrolling iterative
optimization schemes to incrementally optimize the sparse rate reduction objective. In particular, our contributions and outline of the paper are as follows:
- In Section 2.2 we show, using an idealized model for the token distribution, that if one iteratively denoises the tokens towards a family of low-dimensional subspaces, the associated score function assumes an explicit form similar to a self-attention operator seen in transformers.
- In Section 2.3 we derive the multi-head self-attention layer as an unrolled gradient descent step to minimize the lossy coding rate part of the rate reduction, showing another interpretation of the self-attention layer as compressing the token representation.
- In Section 2.4 we show that the multi-layer perceptron which immediately follows the multi-head self-attention in transformer blocks can be interpreted as (and replaced by) a layer which incrementally optimizes the remaining part of the sparse rate reduction objective by constructing a sparse coding of the token representations.
- In Section 2.5 we use this understanding to create a new white-box (fully mathematically interpretable) transformer architecture called CRATE (i.e., Coding RAte reduction TransformEr), where each layer performs a single step of an alternating minimization algorithm to optimize the sparse rate reduction objective.
Hence, within our framework, the learning objective function, the deep learning architecture, and the final learned representation all become white boxes that are fully mathematically interpretable. As the experiments in Section 3 show, the CRATE networks, despite being simple, can already learn the desired compressed and sparse representations on large-scale real-world datasets and achieve performance on par with much more heavily engineered transformer networks (such as ViT) on a wide variety of tasks (e.g., classification and transfer learning).
2 Technical Approach and Justification
2.1 Objective and Approach
We consider a general learning setup associated with real-world signals. We have some random variable  $\mathbf{X} = [\pmb{x}_1, \dots, \pmb{x}_N] \in \mathbb{R}^{D \times N}$  which is our data source; each  $\pmb{x}_i \in \mathbb{R}^D$  is interpreted as a token<sup>1</sup>, and the  $\pmb{x}_i$ 's may have arbitrary correlation structures. We use  $\mathbf{Z} = [\pmb{z}_1, \dots, \pmb{z}_N] \in \mathbb{R}^{d \times N}$  to denote the random variable which defines our representations. Each  $\pmb{z}_i \in \mathbb{R}^d$  is the representation of the corresponding token  $\pmb{x}_i$ . We are given  $B \geq 1$  i.i.d. samples  $X_1, \dots, X_B \sim \mathbf{X}$ , whose tokens are  $\pmb{x}_{i,b}$ . The representations of our samples are denoted  $Z_1, \dots, Z_B \sim Z$ , and those of our tokens are  $z_{i,b}$ . Finally, for a given network, we use  $Z^\ell$  to denote the output of the first  $\ell$  layers when given  $\mathbf{X}$  as input. Correspondingly, the sample outputs are  $Z_i^\ell$  and the token outputs are  $z_{i,b}^\ell$ .
Objective for learning a structured and compact representation. Following the framework of rate reduction [55], we contend that the goal of representation learning is to find a feature mapping  $f\colon \mathbf{X}\in \mathbb{R}^{D\times N}\to \mathbf{Z}\in \mathbb{R}^{d\times N}$  which transforms input data  $\mathbf{X}\in \mathbb{R}^{D\times N}$  with a potentially nonlinear and multi-modal distribution to a (piecewise) linearized and compact feature representation  $\mathbf{Z}\in \mathbb{R}^{d\times N}$ . While the joint distribution of tokens  $(z_{i})_{i = 1}^{N}$  in  $\mathbf{Z}$  may be sophisticated (and task-specific), we further contend that it is reasonable and practical to require that the target marginal distribution of individual tokens  $z_{i}$  should be highly compressed and structured, amenable for compact coding. Particularly, we require the distribution to be a mixture of low-dimensional (say  $K$ ) Gaussian distributions, such that the  $k^{\text{th}}$  Gaussian has mean  $\mathbf{0}\in \mathbb{R}^{d}$ , covariance  $\Sigma_{k}\succeq \mathbf{0}\in \mathbb{R}^{d\times d}$ , and support spanned by the orthonormal basis  $U_{k}\in \mathbb{R}^{d\times p}$ . We denote  $U_{[K]} = (U_k)_{k = 1}^K$  to be the set of bases of all Gaussians. Hence to maximize the information gain [62] for the final token representation, we wish to maximize the rate reduction [6, 46] of the tokens, i.e.,  $\max_{\mathbf{Z}}\Delta R(\mathbf{Z};\mathbf{U}_{[K]}) = R(\mathbf{Z}) - R^{c}(\mathbf{Z};\mathbf{U}_{[K]}),$  where  $R$  and  $R^c$  are estimates of lossy coding rates to be formally defined in (7) and (8). This also promotes token representations  $z_{i}$  from different Gaussians to be incoherent [46]. Since rate reduction is an intrinsic measure of goodness for the representation, it is invariant to arbitrary rotations of the representations. Therefore, to ensure the final representations are amenable to more compact coding, we would like to transform the representations (and their supporting subspaces) so that they become sparse with respect to the standard coordinates of the resulting
For language transformers, tokens roughly correspond to words [28], while for vision transformers, tokens correspond to image patches [40].
representation space. $^2$  The combined rate reduction and sparsification process is illustrated in Figure 1. Computationally, we may combine the above two goals into a unified objective for optimization:
$$
\max  _ {f \in \mathcal {F}} \mathbb {E} _ {\boldsymbol {Z}} \left[ \Delta R (\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}) - \lambda \| \boldsymbol {Z} \| _ {0} \right] = \max  _ {f \in \mathcal {F}} \mathbb {E} _ {\boldsymbol {Z}} \left[ R (\boldsymbol {Z}) - R ^ {c} (\boldsymbol {Z}; \boldsymbol {U} _ {[ K ]}) - \lambda \| \boldsymbol {Z} \| _ {0} \right] \text {s . t .} \boldsymbol {Z} = f (\boldsymbol {X}), (1)
$$
where the  $\ell^0$  norm  $\| Z\| _0$  promotes the sparsity of the final token representations  $Z = f(X)$ .<sup>3</sup> We call this objective "sparse rate reduction."
White-box deep architecture as unrolled incremental optimization. Although easy to state, each term of the above objective can be computationally very challenging to optimize [55, 70]. Hence it is natural to take an approximation approach that realizes the global transformation  $f$  optimizing (1) through a concatenation of multiple, say  $L$ , simple incremental and local operations  $f^{\ell}$  that push the representation distribution towards the desired parsimonious model distribution:
$$
f \colon \boldsymbol {X} \xrightarrow {f ^ {0}} \boldsymbol {Z} ^ {0} \rightarrow \dots \rightarrow \boldsymbol {Z} ^ {\ell} \xrightarrow {f ^ {\ell}} \boldsymbol {Z} ^ {\ell + 1} \rightarrow \dots \rightarrow \boldsymbol {Z} ^ {L} = \boldsymbol {Z}, \tag {2}
$$
where  $f^0: \mathbb{R}^D \to \mathbb{R}^d$  is the pre-processing mapping that transforms input tokens  $x_i \in \mathbb{R}^D$  to their token representations  $z_i^1 \in \mathbb{R}^d$ .
Each incremental forward mapping  $Z^{\ell + 1} = f^{\ell}(Z^{\ell})$ , or a "layer", transforms the token distribution to optimize the above sparse rate reduction objective (1), conditioned on the distribution of its input tokens  $Z^{\ell}$ . In contrast to other unrolled optimization approaches such as the ReduNet [55], we explicitly model the distribution of  $Z^{\ell}$  at each layer, say as a mixture of linear subspaces or sparsely generated from a dictionary. The model parameters are learned from data (say via backward propagation with end-to-end training). This separation of forward "optimization" and backward "learning" clarifies the mathematical role of each layer as an operator transforming the distribution of its input, whereas the input distribution is in turn modeled (and subsequently learned) by the parameters of the layer.
We show that we can derive these incremental, local operations through an unrolled optimization perspective to achieve (1) through Sections 2.3 to 2.5. Once we decide on using an incremental approach to optimizing (1), there are a variety of possible choices to achieve the optimization. Given a model for  $Z^{\ell}$ , say a mixture of subspaces  $U_{[K]}$ , we opt for a two-step alternating minimization process with a strong conceptual basis: first in Section 2.3, we compress the tokens  $Z^{\ell}$  via a gradient step to minimize the coding rate term  $\min_Z R^c(Z; U_{[K]})$ ; second, in Section 2.4, we sparsify the compressed tokens, with a suitably-relaxed proximal gradient step on the difference of the sparsity penalty and the expansion term, i.e.,  $\min_Z [\lambda \| Z \|_0 - R(Z)]$ . Both actions are applied incrementally and repeatedly, as each  $f^{\ell}$  in (2) is instantiated with these two steps.
2.2 Self-Attention via Denoising Tokens Towards Multiple Subspaces
There are many different ways to optimize the objective (1) incrementally. In this work, we propose arguably the most basic scheme. To help clarify the intuition behind our derivation and approximation, in this section (and Appendix A.1) we study a largely idealized model which nevertheless captures the essence of nearly the whole process and particularly reveals the reason why self-attention-like operators arise in many contexts. Assume that  $N = 1$ , and the single token  $\pmb{x}$  is drawn i.i.d. from an unknown mixture of Gaussians  $(\mathcal{N}(\mathbf{0},\pmb{\Sigma}_k))_{k=1}^K$  supported on low-dimensional subspaces with orthonormal bases  $U_{[K]} = (U_k)_{k=1}^K$  and corrupted with additive Gaussian noise  $\pmb{w} \sim \mathcal{N}(\pmb{0},\pmb{I})$ , i.e.,
$$
\boldsymbol {x} = \boldsymbol {z} + \sigma \boldsymbol {w}, \tag {3}
$$
where  $z$  is distributed according to the mixture. Our goal is simply to transform the distribution of the noisy token  $x$  to the mixture of low-dimensional Gaussians  $z$ . Towards incremental construction of a representation  $f$  for this model following (2), we reason inductively: if  $z^{\ell}$  is a noisy token (3) at noise level  $\sigma^{\ell}$ , it is natural to produce  $z^{\ell + 1}$  by denoising at the level  $\sigma^{\ell}$ . In the mean-square sense, the optimal estimate is  $\mathbb{E}[z \mid z^{\ell}]$ , which has a variational characterization (e.g. [12]):
$$
\mathbb {E} [ \boldsymbol {z} | \cdot ] = \arg \min  _ {f} \mathbb {E} _ {\boldsymbol {z}, \boldsymbol {w}} \left[ \| f (\boldsymbol {z} + \sigma^ {\ell} \boldsymbol {w}) - \boldsymbol {z} \| _ {2} ^ {2} \right]. \tag {4}
$$
2That is, having the fewest nonzero entries.
3To simplify the notation, we will discuss the objective for one sample  $X$  at a time with the understanding that we always mean to optimize the expectation.
Setting  $z^{\ell + 1} = \mathbb{E}[z \mid z^{\ell}]$ , (4) thus characterizes the next stage of (2) in terms of an optimization objective based on a local signal model for  $z^{\ell}$ . Moreover, letting  $x \mapsto q^{\ell}(x)$  denote the density of  $z^{\ell}$ , Tweedie's formula [13] allows us to express the optimal representation solving (4) in closed-form:
$$
\boldsymbol {z} ^ {\ell + 1} = \boldsymbol {z} ^ {\ell} + \left(\sigma^ {\ell}\right) ^ {2} \nabla_ {\boldsymbol {x}} \log q ^ {\ell} \left(\boldsymbol {z} ^ {\ell}\right). \tag {5}
$$
Tweedie's formula expresses the optimal representation in terms of an additive correction (in general a nonlinear function of  $z^{\ell}$ ) to the noisy observations by the gradient of the log-likelihood of the distribution of the noisy observations, giving the optimal representation a clear interpretation as an incremental perturbation to the current noisy distribution  $q^{\ell}$ . This connection is well-known in the areas of estimation theory and inverse problems [1, 13, 14, 19, 20, 27, 42], and more recently has found powerful applications in the training of generative models for natural images [4, 15, 22, 43, 44]. Here, we can calculate a closed-form expression for this score function  $\nabla_{\mathbf{x}}\log q^{\ell}$ , which, when combined with (5) and some technical assumptions<sup>4</sup>, gives the following approximation (shown in Appendix A.1). Let  $\otimes$  denote the Kronecker product; then we have
$$
\boldsymbol {z} ^ {\ell + 1} \approx [ \boldsymbol {U} _ {1}, \dots , \boldsymbol {U} _ {K} ] \left[ \operatorname {d i a g} \left(\operatorname {s o f t m a x} \left(\frac {1}{2 \left(\sigma^ {\ell}\right) ^ {2}} \left[ \begin{array}{c} \| \boldsymbol {U} _ {1} ^ {*} \boldsymbol {z} ^ {\ell} \| _ {2} ^ {2} \\ \vdots \\ \| \boldsymbol {U} _ {K} ^ {*} \boldsymbol {z} ^ {\ell} \| _ {2} ^ {2} \end{array} \right]\right)\right) \otimes \boldsymbol {I} _ {p} \right] \left[ \begin{array}{c} \boldsymbol {U} _ {1} ^ {*} \boldsymbol {z} ^ {\ell} \\ \vdots \\ \boldsymbol {U} _ {K} ^ {*} \boldsymbol {z} ^ {\ell} \end{array} \right], \tag {6}
$$
This operation resembles a self-attention layer in a standard transformer architecture with  $K$  heads, sequence length  $N = 1$ , the "query-key-value" constructs being replaced by a single linear projection  $U_k^* z^\ell$  of the token  $z^\ell$ , and the aggregation of head outputs (conventionally modeled by an MLP) done with the two leftmost matrices in (6). We thus derive the following useful interpretation, which we will exploit in the sequel: Gaussian denoising against a mixture of subspaces model leads to self-attention-type layers in the transformation  $f$ . Given an initial sample  $x$  following the model (3), we can repeatedly apply local transformations to the distribution with (6) in order to realize the incremental mapping  $f \colon x \to z$  in (2). These insights will guide us in the design of our white-box transformer architecture in the upcoming subsections.
2.3 Self-Attention via Compressing Token Sets through Optimizing Rate Reduction
In the last subsection, we have seen that the multi-head attention in a transformer resembles the score-matching operator that aims to transform a token  $z^{\ell}$  towards a mixture of subspaces (or degenerate Gaussians). Nevertheless, to carry out such an operation on any data, one needs to first learn or estimate, typically from finite samples, the parameters of the mixture of (degenerate) Gaussians, which is known to be a challenging task [6, 24]. This challenge is made even harder because in a typical learning setting, the given set of tokens are not i.i.d. samples from the mixture of subspaces. The joint distribution among these tokens can encode rich information about the data—for example, co-occurrences between words or object parts in language and image data (resp.)—which we should also learn. Thus, we should compress/denoise/transform such a set of tokens together. To this end, we need a measure of quality, i.e., compactness, for the resulting representation of the set of tokens. A natural measure of the compactness of such a set of tokens is the (lossy) coding rate to encode them up to a certain precision  $\epsilon > 0$  [6, 46]. For a zero-mean Gaussian, this measure takes a closed form. If we view the tokens in  $Z \in \mathbb{R}^{d \times N}$  as drawn from a single zero-mean Gaussian, an estimate of their (lossy) coding rate, subject to quantization precision  $\epsilon > 0$ , is given in [6] as:
$$
R (\boldsymbol {Z}) \doteq \frac {1}{2} \log \det  \left(\boldsymbol {I} + \frac {d}{N \epsilon^ {2}} \boldsymbol {Z} ^ {*} \boldsymbol {Z}\right) = \frac {1}{2} \log \det  \left(\boldsymbol {I} + \frac {d}{N \epsilon^ {2}} \boldsymbol {Z} \boldsymbol {Z} ^ {*}\right). \tag {7}
$$
In practice, the data distribution is typically multi-modal, say an image set consisting of many classes or a collection of image patches as in Figure 1. It is more appropriate to require that the set of tokens map to a mixture of, say  $K$ , subspaces (degenerate Gaussians) [55]. As before we denote the (to be learned) bases of these subspaces as  $U_{[K]} = (U_k)_{k=1}^K$ , where  $U_k \in \mathbb{R}^{d \times p}$ . Although the joint distribution of the tokens  $Z$  is unknown, the desired marginal distribution of each token  $z_i$  is a