[Paper Review] White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?

Yaodong Yu, Sam Buchanan+8 more|arXiv (Cornell University)|Nov 22, 2023

Image and Signal Denoising Methods4 citations

TL;DR

This paper proposes CRATE, a white-box transformer architecture derived from sparse rate reduction—simultaneously optimizing compression and sparsity in representations. By interpreting multi-head self-attention as gradient descent on coding rate and MLP as proximal gradient descent for sparsity, CRATE achieves performance on par with state-of-the-art transformers in vision and language tasks, demonstrating that compression is central to effective representation learning.

ABSTRACT

In this paper, we contend that a natural objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a low-dimensional Gaussian mixture supported on incoherent subspaces. The goodness of such a representation can be evaluated by a principled measure, called sparse rate reduction, that simultaneously maximizes the intrinsic information gain and extrinsic sparsity of the learned representation. From this perspective, popular deep network architectures, including transformers, can be viewed as realizing iterative schemes to optimize this measure. Particularly, we derive a transformer block from alternating optimization on parts of this objective: the multi-head self-attention operator compresses the representation by implementing an approximate gradient descent step on the coding rate of the features, and the subsequent multi-layer perceptron sparsifies the features. This leads to a family of white-box transformer-like deep network architectures, named CRATE, which are mathematically fully interpretable. We show, by way of a novel connection between denoising and compression, that the inverse to the aforementioned compressive encoding can be realized by the same class of CRATE architectures. Thus, the so-derived white-box architectures are universal to both encoders and decoders. Experiments show that these networks, despite their simplicity, indeed learn to compress and sparsify representations of large-scale real-world image and text datasets, and achieve performance very close to highly engineered transformer-based models: ViT, MAE, DINO, BERT, and GPT2. We believe the proposed computational framework demonstrates great potential in bridging the gap between theory and practice of deep learning, from a unified perspective of data compression. Code is available at: https://ma-lab-berkeley.github.io/CRATE .

Motivation & Objective

• To establish a principled, theory-driven framework for representation learning based on data compression and sparsity.
• To explain the success of transformers through the lens of iterative optimization on a unified objective: sparse rate reduction.
• To derive a fully interpretable, mathematically grounded deep network architecture (CRATE) from first principles.
• To unify encoder and decoder functions in a single architecture via invertible, structured denoising and diffusion.
• To empirically validate that CRATE learns meaningful, semantically structured representations across vision and language tasks.

Proposed method

• Formalize representation learning as optimizing sparse rate reduction—a measure combining intrinsic information gain and extrinsic sparsity.
• Derive multi-head self-attention as an approximate gradient descent step on the coding rate of input tokens.
• Model the multi-layer perceptron as a proximal gradient descent step for sparse coding of features.
• Construct CRATE as a stack of alternating attention and sparse coding blocks, ensuring mathematical interpretability.
• Enable invertible decoding by reversing the compression process through structured denoising-diffusion using the same CRATE blocks.
• Implement CRATE in PyTorch with modular components: MSSA (multi-head self-attention) and ISTA (iterative shrinkage thresholding for sparsity).

Experimental results

Research questions

• RQ1 Can the behavior of transformers be explained through a unified optimization objective rooted in information theory and sparsity?
• RQ2 To what extent can a deep network be designed as a white-box, interpretable system by directly deriving it from a principled compression objective?
• RQ3 Does a network derived from sparse rate reduction achieve performance competitive with highly engineered transformer models?
• RQ4 Can the same architecture serve as both encoder and decoder through invertible, structured denoising?
• RQ5 Do the attention maps and learned features in CRATE exhibit emergent semantic structure, such as object-level segmentation?

Key findings

• CRATE achieves performance within 1-2% of state-of-the-art models like ViT, MAE, DINO, BERT, and GPT2 on image and text benchmarks.
• The attention maps in CRATE layers exhibit emergent semantic segmentation, with high attention concentration on object boundaries and salient features.
• Layer-wise analysis confirms that CRATE learns hierarchical, structured representations, with early layers capturing low-level patterns and deeper layers capturing high-level semantics.
• The inverse of the CRATE encoder can be realized via structured denoising-diffusion, demonstrating that the same architecture supports both encoding and decoding.
• The model's performance remains stable across ablation studies, confirming the robustness of the sparse rate reduction objective.
• The learned representations are parsimonious and structured, with features concentrated on incoherent subspaces, aligning with the theoretical compression goal.