PaperlayerView exact revision on arXiv

Source-linked AI summary

White-Box Transformers via Sparse Rate Reduction

Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin D. Haeffele, Yi Ma

arXiv:2306.01129v1cs.LG

TL;DR

Transformers have achieved success but lack a rigorous mathematical interpretation of how data distributions relate to their learned representations. This paper derives transformer-like networks as incremental optimizers of sparse rate reduction, with attention compressing token sets and the MLP sparsifying them. The resulting interpretable networks perform on par with more heavily engineered transformers on a variety of tasks.

  • Problem

    Transformers are empirically designed, and the statistical and geometric relationship between data distributions and their learned representations remains unclear.

    Source: §1 · p. 2

  • Method

    The paper derives transformer-like layers as alternating optimization steps for sparse rate reduction, with attention compressing token sets and the MLP constructing sparse token codes.

    Source: §1 · pp. 2–3

  • Results

    CRATE networks achieve performance on par with more heavily engineered transformers such as ViT across tasks including classification and transfer learning.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The derived architectures are mathematically interpretable and consistent layer by layer with their design objective.

    Source: §4 · p. 10

    The objective assumes the target marginal distribution of individual tokens is a mixture of low-dimensional Gaussian distributions.

    Source: §2.1 · p. 3

Abstract

from arXiv · show

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 \url{https://github.com/Ma-Lab-Berkeley/CRATE}.

1 Introduction

The introduction proposes sparse rate reduction as a framework for deriving interpretable transformer-like networks that incrementally optimize compressed and sparse representations. It frames this approach against the limited mathematical interpretation of transformers and reports performance on par with more heavily engineered transformers.

1 Introduction

  • Transformers lack a rigorous mathematical interpretation, while diffusion models use black-box denoising and do not establish a clear correspondence between initial features and data samples.The paper says diffusion models therefore do not provide a parsimonious or interpretable representation of the data distribution.Source: §1 · p. 2
  • Explicit representation-learning approaches can be interpretable, but narrowly defined desired properties may make good performance on large real-world datasets difficult to achieve.These approaches include model-based sparse coding and model-free methods such as contrastive learning and rate reduction.Source: §1 · p. 2
  • The framework derives transformer-like layers as incremental optimizers of a unified sparse rate reduction objective for compressed and sparse token representations.The authors present the resulting architecture as mathematically interpretable.Source: §1 · pp. 2–3
  • Under an idealized token-distribution model, denoising toward low-dimensional subspaces yields a score function similar to transformer self-attention, which is also derived as a gradient step for lossy-rate compression.The attention interpretation concerns compression of token representations.Source: §1 · p. 3
  • The transformer’s subsequent MLP can be interpreted as incrementally optimizing the objective’s remaining part by constructing sparse codes for token representations.Source: §1 · p. 3
  • CRATE achieves performance on par with more heavily engineered transformers such as ViT across tasks including classification and transfer learning.The authors report this despite CRATE's simplicity.Source: §1 · p. 3

2 Technical Approach and Justification

CRATE frames representation learning as incremental optimization of sparse rate reduction, alternating token-set compression and sparsification to produce compact representations through mathematically interpretable layers.

2.1 Objective and Approach

  • Sparse rate reduction targets token representations that maximize the difference between overall and subspace-conditioned coding rates, while promoting sparsity in standard coordinates.The target distribution is a mixture of low-dimensional Gaussians supported on subspaces.Source: §2.1 · pp. 3–4

2.2 Self-Attention via Denoising Tokens Towards Multiple Subspaces

  • Under an idealized single-token Gaussian-mixture model, denoising via the score of the noisy-token distribution yields an operation resembling self-attention.The optimal next representation is a conditional expectation, expressed using Tweedie’s formula as a score-gradient update.Source: §2.2 · pp. 4–5

2.3 Self-Attention via Compressing Token Sets through Optimizing Rate Reduction

  • For token sets, a coding-rate gradient step is approximated by MSSA, which shares subspace-basis projections across query, key, and value and aggregates head outputs.With one token, an aggressive step, and tuned quantization error, MSSA becomes the ideal denoiser apart from linear rather than nonlinear head aggregation.Source: §2.3 · p. 6

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

  • A global approximately orthogonal dictionary supports sparsification of compressed tokens; the approach preserves coding rate approximately and avoids a poorly scalable matrix-inverse gradient.The sparse representation program is relaxed to LASSO and incrementally optimized with an ISTA step.Source: §2.4 · p. 7

2.5 The Overall White-Box CRATE Architecture

  • CRATE layers combine MSSA compression with ISTA sparsification, using layer-specific subspace bases and dictionaries learned from data via backpropagation.The resulting architecture transforms tokens toward a compact, sparse union of incoherent subspaces.Source: §2.5 · pp. 7–8

3 Experiments

CRATE is evaluated both by checking whether its layers meet their design objectives and by measuring performance on ImageNet-1K and transfer tasks. The experiments report progressive compression and sparsification, alongside performance similar to ViT at comparable model size.

3 Experiments

  • Experiments assess whether CRATE’s layers meet their design goals and test performance on ImageNet-1K and fine-tuned downstream datasets.The models vary in token dimension, number of heads, and number of layers; downstream datasets include CIFAR10/100, Oxford Flowers, and Oxford-IIT-Pets.Source: §3 · p. 8

3.1 In-depth Layer-wise Analysis of CRATE

  • Across most layers, MSSA outputs show improved compression and ISTA outputs improved sparsity as layer depth increases.The layer-wise measures are evaluated on ImageNet-1K samples; the final layer’s sparsity measure rises because of the added classification layer.Source: §3.1 · p. 9, Figure 3 · p. 9
  • Untrained CRATE does not optimize the objectives effectively; during training, learned subspaces and dictionaries help the designed blocks optimize them more effectively.Figure 4 compares compression and sparsification terms across training durations.Source: §3.1 · p. 9, Figure 4 · p. 9
  • ISTA outputs become sparser and more low-rank across layers, bringing token representations closer to linear subspaces.The authors suggest the last-layer sparsity increase may reflect features shared across object classes.Source: §3.1 · pp. 9–10

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

  • With similar parameter counts, CRATE achieves similar ImageNet-1K and transfer-learning performance to ViT, while scaling CRATE consistently improves performance.The comparison uses CRATE-Base (22.08 million parameters) and ViT-Small (22.05 million parameters).Source: Table 1 · p. 10, §3.2 · p. 10

3.1 In-depth Layer-wise Analysis of CRATE

  • Learned subspaces are approximately incoherent, and their incoherence increases at later layers, suggesting token representations become more separable.The authors note this pattern mirrors that in other popular deep networks.Source: §3.1 · p. 10

4 Conclusion

The paper presents CRATE as an interpretable transformer-like architecture derived from incremental optimization of compressed and sparse representations. It reports performance close to seasoned transformers and frames the framework as a source of theoretical guidance for architecture design.

4 Conclusion

  • The framework derives deep transformer-like architectures as incremental optimization schemes for learning compressed and sparse representations.The authors describe the resulting architectures as mathematically interpretable and consistent layer by layer with their design objective.Source: §4 · p. 10
  • Despite their simple design, the networks achieve performance close to seasoned transformers on large-scale real-world datasets and tasks.Source: §4 · p. 10
  • The authors argue the framework provides theoretical guidelines for designing and justifying new deep architectures for representation learning.Source: §4 · p. 10

A.1 Companion to Section 2.2

The companion section revisits the technical contributions of the denoising analysis and supplies proofs. Its idealized argument concerns single-token Gaussian mixtures, and one approximation is presented with a possible small-noise asymptotic interpretation.

A.1 Companion to Section 2.2

  • The appendix supplies technical proofs for the denoising contribution, using notation that maps the variables to layer-indexed representations and noise.Source: §A.1 · p. 15
  • The analysis presents a single-token Gaussian-mixture denoising mechanism and relates its representation-learning scheme to self-attention at sequence length one.It interprets self-attention as a generalized mechanism for denoising a mixture-of-Gaussian marginal model for a token set.Source: §A.1 · p. 15
  • Approximation 2 may admit asymptotic equivalence as σ approaches zero, particularly below the smallest nonzero covariance eigenvalue; the softmax limit is the delicate step.The authors present this as a possibility rather than a proved equivalence.Source: §A.1 · p. 18

A.2 Companion to Section 2.3

The section develops an unsupervised coding-rate formulation over subspaces and derives self-attention as an approximate gradient step that compresses token representations. Its derivation uses approximations and technical assumptions, and the authors do not claim the resulting network is optimally designed.

A.2 Companion to Section 2.3

  • The generalized coding-rate function measures coding rate relative to subspaces rather than classes, making the formulation unsupervised.Source: §A.2 · p. 18
  • A gradient step on the coding-rate objective nearly recovers multi-head self-attention, with the query, key, and value operators sharing one projection.The projection is interpreted as projection onto a single subspace.Source: §A.2 · p. 18
  • The authors use simple approximations and technical assumptions and do not claim their network is optimally designed.They present compression, denoising, sparsification, and unrolled optimization as principles for future interpretable architectures.Source: §A.2 · p. 18
  • For computational efficiency, the derivation approximates an expensive matrix inverse and uses softmax to focus regression and compression on tokens assigned to the same subspace.The inverse is approximated using the first-order von Neumann expansion.Source: §A.2 · p. 19

A.3 Companion to Section 2.4

The section analyzes the sparsifying step of sparse rate reduction, connecting it to LASSO and transformer MLPs, then deriving an ISTA-like update with proximal majorization-minimization. The derivation depends on simplifying assumptions, and the paper identifies broader dictionary choices as future work.

A.3 Companion to Section 2.4

  • The paper links the second alternating step heuristically to LASSO and its unrolled proximal-gradient step to the transformer MLP; the appendix derives a more rigorous alternative.The alternative uses majorization-minimization and produces an update resembling ISTA.Source: §A.3 · p. 20
  • With normalized, orthogonal post-compression token columns, incremental minimization yields a slightly simplified version of the ISTA block.These assumptions provide a mathematical justification for the earlier ISTA derivation.Source: §A.3 · p. 20
  • For an orthogonal dictionary, the update resembles ISTA and is interpreted as sparsifying compressed tokens in a learnable dictionary.This version lacks a skip connection.Source: §A.3 · p. 22
  • The derivation treats orthogonal dictionaries; overcomplete dictionaries are left for future work and are expected to improve CRATE scalability.The authors also describe paired dictionaries and near-orthogonal dictionaries as possible extensions.Source: §A.3 · p. 22

A.3.1 Auxiliary Lemmas

  • An auxiliary SVD estimate bounds σ/(1+ασ^2) by 1/(2*sqrt(α)), supporting the stated supremum bound.Source: §A.3.1 · p. 24

B Additional Experiments and Details

The additional experiments report strong CRATE performance on large-scale datasets, including ImageNet-1K classification, and strong transfer learning. They also emphasize layer-wise interpretability and report that studied architecture modifications do not significantly hurt empirical performance.

B Additional Experiments and Details

B.1 Implementation details

  • The appendix provides additional implementation details for applying CRATE to vision tasks.Source: §B.1 · p. 25

B.1.1 Architecture of CRATE

For implementation simplicity, the CRATE architecture changes the compression step by replacing its projection term with a trainable parameter. The paper also supplies PyTorch-style code and algorithms for the overall architecture and its transformer, self-attention, and MLP blocks.

B.1.1 Architecture of CRATE

  • For implementation simplicity, the compression step replaces its projection term with a trainable parameter W in R^(d×pK).Source: §B.1.1 · p. 25
  • The replacement parameter W takes the place of the term p Nϵ^2 [U1, . . . , UK] in the MSSA operator.Source: §B.1.1 · p. 25
  • PyTorch-style code gives the overall architecture and algorithms detailing the transformer, MSSA, and ISTA blocks.Source: §B.1.1 · p. 25

B.1.2 Training Setup

The appendix specifies ImageNet pre-training and fine-tuning settings and provides pseudocode for the CRATE classifier and transformer components. It also introduces additional experiments, including layer-wise measurements, visualizations, and ablations.

B.1.2 Training Setup

  • ImageNet-1K pre-training uses Lion for 150 epochs, with a 5-epoch linear warm-up, cosine decay, learning rate 2.4 × 10−4, weight decay 0.5, and batch size 2,048.Augmentation is limited to random cropping and horizontal flipping, with label smoothing 0.1; one CRATE-Base epoch takes about 240 seconds on 16 A100 40GB GPUs.Source: §B.1.2 · p. 25
  • Fine-tuning uses AdamW with learning rate 5 × 10−5, weight decay 0.01, and batch size 512 across CIFAR10/100, Oxford Flowers-102, Oxford-IIIT-Pets, and ImageNet Real evaluation.Inputs are resized to 224; augmentation includes random cropping, horizontal flipping, and random augmentation with n = 2 and m = 14.Source: §B.1.2 · p. 25
  • The classifier pseudocode embeds image patches, adds a class token and positional embeddings, applies the transformer, pools the tokens, and sends the result to a normalized linear classification head.The listed forward pass uses either mean pooling or the class token.Source: §B.1.2 · p. 26

B.2 Experimental Results

  • The appendix introduces additional CRATE experiments covering layer-wise measurements, visualizations, and ablation studies.The provided pseudocode also presents attention and feedforward components for the CRATE transformer block.Source: §B.1.2 · pp. 26–27, §B.2 · p. 28

B.2.1 Layer-wise Evaluation and Visualization

Across CRATE-Tiny, Base, and Large, the compression and sparsity terms improve with layer depth. Layer-wise visualizations show increasingly sparse, structured token representations and include additional token and subspace views.

B.2.1 Layer-wise Evaluation and Visualization

  • Both the compression term at MSSA outputs and sparsity at ISTA outputs improve as layer index increases across CRATE-Tiny, Base, and Large.Figure 5 plots the compression term Rc(Zℓ+1/2) and normalized sparsity ∥Zℓ+1∥0/(d · N) by layer.Source: §B.2.1 · p. 28, Figure 5 · p. 28
  • Figure 6 presents CRATE-Tiny token representations Zℓ for each of layers 1 through 12, displaying a randomly extracted 50×50 submatrix from each layer’s representation for visual clarity.Source: §B.2.1 · p. 29, Figure 6 · p. 29
  • Figure 7 visualizes multi-head subspaces across layers; four directions are selected from each subspace for display, and the subspaces are described as incoherent to capture different token-set views.Source: §B.2.1 · pp. 29–30, Figure 7 · p. 30

B.2.2 Additional Layer-wise Visualization

B.3 CRATE Ablation

The ablation evaluates CRATE’s Top 1 accuracy under varied training and ISTA-block hyperparameters, and reports consistently satisfactory performance across a diverse range of settings.

B.3 CRATE Ablation

  • CRATE maintains consistently satisfactory performance across the evaluated hyperparameter settings.Table 2 reports Top 1 accuracy for CRATE variants trained on ImageNet.Source: §B.3 · p. 33, Table 2 · p. 33
  • The study varies the number of epochs, weight decay, and learning rate to evaluate CRATE’s performance across training settings.Source: §B.3 · p. 33, Table 2 · p. 33
  • The ablation also varies the ISTA-block step size η and regularization term λ.Source: §B.3 · p. 33, Table 2 · p. 33

B.4 Exploring Architecture Variants

The paper evaluates alternative attention and sparsification mechanisms, then reports that CRATE maintains reasonable ImageNet-1K performance after additional architectural simplifications.

B.4 Exploring Architecture Variants

  • The authors describe these design choices as modular and principled, while noting that more sophisticated analysis could yield more complicated architectures with better practical performance.Source: §B.4 · p. 33
  • The explored variants directly implement the compression-inspired attention mechanism or replace the ISTA sparsification step with a majorization-minimization proximal step.Source: §B.4 · p. 33
  • After additional architectural constraints and simplifications, CRATE maintains reasonable performance on ImageNet-1K.Table 3 reports Top 1 accuracy for CRATE architecture variants trained on ImageNet.Source: §B.4 · p. 33, Table 3 · p. 33
Loading 2306.01129v1…