An In-depth Investigation of Sparse Rate Reduction in Transformer-like Models

TL;DR

Surprisingly, it is found that Sparse Rate Reduction has a positive correlation coefficient and outperforms other baseline measures, such as path-norm and sharpness-based ones, and generalization can be improved using SRR as regularization on benchmark image classification datasets.

Abstract

Deep neural networks have long been criticized for being black-box. To unveil the inner workings of modern neural architectures, a recent work \cite{yu2024white} proposed an information-theoretic objective function called Sparse Rate Reduction (SRR) and interpreted its unrolled optimization as a Transformer-like model called Coding Rate Reduction Transformer (CRATE). However, the focus of the study was primarily on the basic implementation, and whether this objective is optimized in practice and its causal relationship to generalization remain elusive. Going beyond this study, we derive different implementations by analyzing layer-wise behaviors of CRATE, both theoretically and empirically. To reveal the predictive power of SRR on generalization, we collect a set of model variants induced by varied implementations and hyperparameters and evaluate SRR as a complexity measure based on its correlation with generalization. Surprisingly, we find out that SRR has a positive correlation coefficient and outperforms other baseline measures, such as path-norm and sharpness-based ones. Furthermore, we show that generalization can be improved using SRR as regularization on benchmark image classification datasets. We hope this paper can shed light on leveraging SRR to design principled models and study their generalization ability.

Figures (6)

• Figure 1: In a simplified attention-only experiment, MSSA operator with skip connection actually implements an ascent method on $R^c(\boldsymbol{Z};\boldsymbol{U})$, opposed to its design purpose (left). This is due to an artifact in approximation with its second-order term. (right)
• Figure 2: Sparse rate reduction measure $\lambda \|\boldsymbol{Z}\|_0 + R^c(\boldsymbol{Z};\boldsymbol{U})- R(\boldsymbol{Z})$ of CRATE and its variants evaluated at different layers and epochs on CIFAR-10.
• Figure 3: Sparse rate reduction measure $\lambda \|\boldsymbol{Z}\|_0 + R^c(\boldsymbol{Z};\boldsymbol{U})- R(\boldsymbol{Z})$ of CRATE and its variants evaluated at different layers and epochs on CIFAR-100.
• Figure 4: A scatter plot illustrating the value of SRR measure and generalization gap across CRATE and variants with network width $d=384$.
• Figure 5: (a) Original gradient update, i.e, \ref{['eq:original']}. (b) Update from second-order Taylor expansion, i.e., \ref{['eq:omission']}. (c) Update from removing the second-order term from \ref{['eq:omission']}. (d) Update from removing the first-order term from \ref{['eq:omission']}. (e) Update from further adding softmax, i.e., \ref{['eq:approx descent']}.