White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is? (2311.13110v4)

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 .

• The paper introduces a unified sparse rate reduction objective to derive an interpretable, compression-driven transformer architecture.
• It employs unrolled optimization via multi-head subspace self-attention and ISTA for iterative encoding and sparsification.
• Experimental results across vision and text tasks confirm competitive performance and enhanced model interpretability.

White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?

The paper "White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?" presents a principled approach to understanding and designing deep network architectures, particularly transformers, based on the concept of representation learning through structured lossy compression. The authors introduce a framework that integrates the principles of sparsity and rate reduction to derive a fully interpretable transformer-like architecture. This essay examines the paper's methodology, results, and implications in the field of large-scale representation learning.

Sparse Rate Reduction: A Unified Objective

The central hypothesis of the paper is that an effective representation learning framework can be derived by optimizing an objective that combines rate reduction and sparsity—termed as sparse rate reduction. The idea is to map a high-dimensional, complex data distribution into a lower-dimensional, compact, and structured representation space. This objective aims to balance intrinsic information gain, measured by rate reduction, and extrinsic structural simplicity, achieved through sparsity. The paper's authors demonstrate this objective's efficacy in simultaneously promoting compression, linearization, and sparsity in learned representations.

Deriving Transformer-Like Architectures

Building upon the defined objective, the authors employ unrolled optimization strategies to incrementally encode the data distribution into the desired parsimonious structure. Specifically, the representation is refined iteratively through layers that conduct alternating steps of compression and sparsification. Compression is realized via a multi-head subspace self-attention (MSSA) operator, derived from the gradient descent on the coding rate term $R^c$. Sparsification, on the other hand, uses an iterative shrinkage-thresholding algorithm (ISTA) to effectively introduce axis-aligned structures into the feature space.

These iterative steps are systematically integrated into a network architecture named crate (Coding-RATE transformer), which leverages mathematically coherent operators to compress and sparsify input data, achieving competitive performance on par with empirically designed black-box counterparts.

Structured Denoising and Diffusion for Decoding

A notable contribution of the paper is its extension to autoencoding by linking compression, diffusion, and denoising processes. By interpreting compression as structured denoising, the work derives a decoder that is a precise qualitative inverse of the encoder. This additional integration implies that both encoding and decoding in the crate framework can be accomplished through the same architectural paradigm, elegantly supporting both discriminative and generative tasks.

Experimental Validation and Interpretability

The empirical evaluation of crate spans supervised and self-supervised learning tasks across vision (e.g., image classification, masked autoencoding) and text (e.g., BERT and GPT-style pretraining). Despite its simplicity and reductive assumptions, crate is shown to achieve remarkable performance, highlighting an avenue where theory meets practice effectively. Moreover, the inherent design transparency allows layer-by-layer analysis, corroborating the intended functional roles of MSSA and ISTA blocks by verifying improvements in compression and sparsity objectives across network layers.

Future Directions and Implications

By demystifying the success of transformers through rigorous optimization-driven design, the paper opens new pathways for developing efficient, white-box learning systems. With this work, future research can explore enhancements in sparsity and rate reduction strategies, extend structured diffusion techniques to more complex generative models, and leverage this framework to build even more scalable and interpretable network architectures.

In summary, this paper provides convincing evidence that effective deep learning models closely resemble optimization procedures advancing towards greater compression and reduction of representation complexity. From a theoretical standpoint, the insights gleaned from this paper support the premise that achieving optimal compression might be the foundation upon which more general forms of intelligence can develop, both in artificial systems and potentially in understanding natural cognitive processes.