Attention-Only White-Box Transformer via LeJEPA-Based Self-Supervised Pretraining

# Attention-Only White-Box Transformer via LeJEPA-Based Self-Supervised Pretraining
Source: [https://arxiv.org/html/2608.04213](https://arxiv.org/html/2608.04213)
###### Abstract

Existing studies on self\-supervised learning for white\-box networks typically decouple the derivation of white\-box networks via optimization algorithms from self\-supervised learning paradigms\. In this work, we instead revisit the two components from a joint perspective\. The LeJEPA\-based self\-supervised framework assumes an isotropic Gaussian distribution as the optimal embedding distribution for downstream tasks, which is conceptually equivalent to the expansion termR​\(Z\)R\(Z\)in the sparse rate reduction objective guiding white\-box Transformer optimization\. Building on this observation, we use the LeJEPA self\-supervised paradigm to optimizeR​\(Z\)R\(Z\), and derive the remaining termsRc​\(Z∣U\[K\]\)\+λ​∥Z∥0R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\+\\lambda\\lVert Z\\rVert\_\{0\}via the alternating direction method of multipliers \(ADMM\) into an attention\-only Transformer that dispenses with the ISTA structure or MLP layers of the original design\. Experimental results demonstrate that our attention\-only white\-box Transformer achieves classification accuracies of88\.88%88\.88\\%on CIFAR\-10 and63\.54%63\.54\\%on CIFAR\-100 at the Base scale under the LeJEPA self\-supervised paradigm, while the original white\-box Transformer CRATE achieves classification accuracies of89\.18%89\.18\\%on CIFAR\-10 and63\.56%63\.56\\%on CIFAR\-100\. Our model achieves competitive performance while reducing the parameter count by roughly31%31\\%\. Beyond the white\-box setting, we further investigate standard ViTs and find that replacing all MLP blocks with ReLU activations under knowledge distillation removes approximately 66% of the parameters while preserving competitive accuracy, motivating further investigation into the potential redundancy of MLP modules in standard ViT architectures\.

## Introduction

Transformers have become a foundational architecture in modern deep learning and have achieved remarkable performance in computer vision, natural language processing, and other fields\(Vaswaniet al\.[2017](https://arxiv.org/html/2608.04213#bib.bib1); Dosovitskiyet al\.[2021](https://arxiv.org/html/2608.04213#bib.bib33)\)\. However, the design of the Transformer architecture and its many variants remains largely empirical and lacks a rigorous mathematical interpretation\(He and Hofmann[2024](https://arxiv.org/html/2608.04213#bib.bib12); Donget al\.[2021](https://arxiv.org/html/2608.04213#bib.bib11); Caronet al\.[2021](https://arxiv.org/html/2608.04213#bib.bib5); Chenet al\.[2021](https://arxiv.org/html/2608.04213#bib.bib10)\)\. This has largely hindered the development of new Transformer variants with improved efficiency or interpretability\.

In recent years, white\-box Transformer architectures, exemplified by the Coding Rate Transformer \(CRATE\), have offered a promising new perspective on addressing the aforementioned problem\. Grounded in the principle of sparse rate reduction, CRATE characterizes an ideal representation through the coordinated interplay among an expansion termR​\(Z\)R\(Z\), a compression termRc​\(Z∣U\[K\]\)R^\{c\}\(Z\\mid U\_\{\[K\]\}\), and a sparsity\-inducing termλ​∥Z∥0\\lambda\\lVert Z\\rVert\_\{0\}\(Yuet al\.[2023a](https://arxiv.org/html/2608.04213#bib.bib21)\)\. In the resulting unrolled architecture, the compression term is approximated by a single gradient\-descent step, giving rise to a Multi\-Head Subspace Self\-Attention \(MSSA\) module\. Meanwhile, the expansion and sparsity terms are integrated through the introduction of an orthogonal sparse dictionaryDDand jointly optimized using an iterative shrinkage\-thresholding algorithm \(ISTA\) module\. This unrolling framework explicitly interprets the computations performed by the core network components as iterative steps for solving well\-defined objective subproblems, thereby providing the Transformer with a clear and unified mathematical interpretation\.

Building on this foundation, white\-box Transformer architectures have continued to evolve in recent years\. CRATE\-α\\alphaimproves both the sparse coding module and the training strategy, substantially enhancing the model’s scalability to large\-scale vision tasks\(Yanget al\.[2024](https://arxiv.org/html/2608.04213#bib.bib22)\)\. The Token Statistics Transformer \(ToST\), meanwhile, derives an interpretable attention mechanism with linear time and memory complexity from a variational formulation of the maximal coding rate reduction objective\(Wuet al\.[2024](https://arxiv.org/html/2608.04213#bib.bib23)\)\. In parallel, this line of research has been extended to a diverse range of tasks, including unsupervised object segmentation, language modeling, and long\-sequence modeling, highlighting the broad potential of white\-box networks in terms of interpretability, computational efficiency, and cross\-domain applicability\(Yuet al\.[2023b](https://arxiv.org/html/2608.04213#bib.bib24); Paiet al\.[2024](https://arxiv.org/html/2608.04213#bib.bib25)\)\.

White\-box networks seek to uncover the low\-dimensional subspace structure underlying token representations directly from data\. Because coding\-rate objectives can be formulated without semantic labels, they provide a possible route to self\-supervised representation learning\(Yuet al\.[2020](https://arxiv.org/html/2608.04213#bib.bib2)\)\. Several recent studies have attempted to bridge this gap\. For instance, CRATE\-MAE incorporates white\-box encoders and decoders into a masked autoencoding framework\(Paiet al\.[2024](https://arxiv.org/html/2608.04213#bib.bib25); Heet al\.[2022](https://arxiv.org/html/2608.04213#bib.bib8)\), while EMP\-SSL introduces the total coding rateR​\(Z\)R\(Z\)into a self\-supervised objective to promote representation expansion\(Tonget al\.[2023](https://arxiv.org/html/2608.04213#bib.bib26)\)\. However, these approaches remain subject to notable limitations: the former relies on a reconstruction loss that lacks a clear theoretical interpretation, whereas the latter retains a black\-box network architecture\. These limitations motivate a fundamental question: how can models be trained under a self\-supervised paradigm while preserving theoretical self\-consistency between the white\-box architecture and its learning objective?

To address this question, we revisit the three components of the white\-box objective: the expansion term is transferred to and optimized through the LeJEPA self\-supervised objective, while the compression and sparsity terms are optimized using the alternating direction method of multipliers \(ADMM\)\(Boydet al\.[2011](https://arxiv.org/html/2608.04213#bib.bib29); Yang and Zhang[2011](https://arxiv.org/html/2608.04213#bib.bib35)\)and unrolled into an attention\-only Transformer architecture\. Our approach builds on the close correspondence, in representation space, between the expansion term in the white\-box objective and the isotropy constraint imposed by LeJEPA\(Balestriero and LeCun[2025](https://arxiv.org/html/2608.04213#bib.bib27); Assranet al\.[2023](https://arxiv.org/html/2608.04213#bib.bib9)\)\. Based on this correspondence, we formulate a new optimization objective by replacing the expansion termR​\(Z\)R\(Z\)with the LeJEPA objective, while retaining the subspace compression termRc​\(Z∣U\[K\]\)R^\{c\}\(Z\\mid U\_\{\[K\]\}\)and the sparsity termλ​∥Z∥0\\lambda\\lVert Z\\rVert\_\{0\}\. We then optimize the resulting objective using ADMM\. Unrolling the resulting ADMM iterations yields a white\-box self\-supervised architecture consisting solely of attention modules\. We refer to the resulting Attention\-Only Transformer derived from ADMM as AoT\-ADMM\. The overall self\-supervised training framework of AoT\-ADMM is illustrated in Figure[1](https://arxiv.org/html/2608.04213#Sx2.F1)\.

Our main contributions are summarized as follows:

- •We propose an attention\-only white\-box Transformer based on the LeJEPA self\-supervised objective\. The model uses the LeJEPA self\-supervised learning paradigm to optimize the expansion termR​\(Z\)R\(Z\)in maximal sparse rate reduction, while the remaining compression and sparsity terms follow the white\-box Transformer derivation paradigm, thereby combining self\-supervised learning with a white\-box optimization objective\.
- •We optimize the expansion termR​\(Z\)R\(Z\)with the LeJEPA self\-supervised learning paradigm and optimize the remaining white\-box objective terms using ADMM\. The compression term corresponds to the attention structure of the Transformer, while the sparsity term admits a closed\-form solution through a ReLU operator, ultimately yielding an attention\-only white\-box Transformer\.
- •Experimental results show that AoT\-ADMM achieves performance comparable to CRATE while reducing parameters by approximately31%31\\%\. At the Base scale, our model reaches88\.88%88\.88\\%on CIFAR\-10 and63\.54%63\.54\\%on CIFAR\-100, compared with CRATE’s89\.18%89\.18\\%and63\.56%63\.56\\%\. It also consistently outperforms AoT at similar parameter counts\. On standard ViTs, replacing all MLP blocks with ReLU activations under knowledge distillation removes approximately66%66\\%of the parameters while preserving competitive accuracy, further demonstrating the potential redundancy of MLP modules\.

## Related Work

### White\-Box Self\-Supervised Learning

Conventional Transformer modules are typically selected through empirical trial and error\(He and Hofmann[2024](https://arxiv.org/html/2608.04213#bib.bib12); Donget al\.[2021](https://arxiv.org/html/2608.04213#bib.bib11)\), without a strict mathematical correspondence between the architecture and its learning objective\. CRATE provides an alternative by deriving its basic operators from sparse rate reduction, such that each network operation has an explicit optimization interpretation\(Yuet al\.[2023a](https://arxiv.org/html/2608.04213#bib.bib21)\)\. This theoretical property makes white\-box models naturally suited to self\-supervised representation learning\. Nevertheless, previous attempts have not established a consistent theoretical source for both the training objective and the network architecture\. CRATE\-MAE adopts a white\-box encoder but still relies on an uninterpretable reconstruction loss\(Paiet al\.[2024](https://arxiv.org/html/2608.04213#bib.bib25)\), whereas EMP\-SSL introduces the total coding rateR​\(Z\)R\(Z\)but continues to employ a black\-box architecture\(Tonget al\.[2023](https://arxiv.org/html/2608.04213#bib.bib26)\)\. In contrast, we first establish the theoretical connection between an isotropic Gaussian prior and the coding\-rate expansion termR​\(Z\)R\(Z\)and use it as the training objective, and then derive the complete network architecture by solving the remaining white\-box objective terms\. To our knowledge, this is the first self\-supervised formulation in which both the training loss and the forward computation originate from a unified white\-box theoretical framework\.

### Attention\-Only Transformers

Transformer simplification studies have shown that some block components can be removed or rearranged\(He and Hofmann[2024](https://arxiv.org/html/2608.04213#bib.bib12)\)\. However, naively removing the feed\-forward branch can weaken representation diversity, and theoretical analysis has identified rank\-collapse behavior in pure\-attention stacks\(Donget al\.[2021](https://arxiv.org/html/2608.04213#bib.bib11)\)\. An interesting study by Wang et al\. provides such a derivation: they model tokens as a noisy mixture of low\-rank Gaussians, treat representation learning as compressing noisy tokens back onto their subspaces, read multi\-head subspace self\-attention as an iterative denoising operator, and unroll that operator into a white\-box Attention\-Only Transformer \(AoT\) built from attention and residual connections alone\(Wanget al\.[2025](https://arxiv.org/html/2608.04213#bib.bib28)\)\. Their target is subspace denoising, so the attention\-only form follows from recovering clean representations rather than from optimizing a full representation\-learning objective\. In contrast, our layer is derived from a different optimization problem: both its attention\-only architecture and its training procedure emerge naturally from the white\-box objective\. Specifically, we preserve the sparse rate reduction objective, move its expansion term into the training loss, and obtain the attention\-only architecture by unrolling the solution to the remaining compression and sparsity terms\.

![Refer to caption](https://arxiv.org/html/2608.04213v1/x1.png)Figure 1:Overview of the proposed white\-box self\-supervised framework\. LeJEPA optimizes the global expansion term, while ADMM unrolls the compression and sparsity terms into the attention\-only AoT\-ADMM encoder, deriving both the training objective and architecture from a unified white\-box objective\.

## Method

This section presents the complete derivation of an attention\-only white\-box Transformer under LeJEPA\-based self\-supervised pretraining\. As illustrated in Figure[1](https://arxiv.org/html/2608.04213#Sx2.F1), the proposed framework derives both the training objective and the network architecture from the sparse rate reduction principle\. We first establish the connection between Sketched Isotropic Gaussian Regularization \(SIGReg\) in LeJEPA and the global expansion termR​\(Z\)R\(Z\), which enables the original objective to be separated into loss\-level and architecture\-level components\. We then apply ADMM to the remaining compression and sparsity terms, yielding a three\-state recursion\. Finally, these ADMM iterations are unrolled along network depth to construct the attention\-only forward architecture\.

### Optimizing the Global Expansion Term

Following maximal coding\-rate reduction and its sparse white\-box Transformer formulation\(Yuet al\.[2020](https://arxiv.org/html/2608.04213#bib.bib2),[2023a](https://arxiv.org/html/2608.04213#bib.bib21)\), sparse rate reduction characterizes a desirable representation as globally expanded, compact within low\-dimensional subspaces, and sparse\. LetZ=f​\(X\)∈ℝd×NZ=f\(X\)\\in\\mathbb\{R\}^\{d\\times N\}denoteNNtoken representations and letU\[K\]=\(Uk\)k=1KU\_\{\[K\]\}=\(U\_\{k\}\)\_\{k=1\}^\{K\}denote a collection of low\-dimensional subspace bases\. The corresponding objective is

minf∈ℱ⁡𝔼Z=f​\(X\)​\[Rc​\(Z∣U\[K\]\)−R​\(Z\)\+λ​∥Z∥0\],\\min\_\{f\\in\\mathcal\{F\}\}\\;\\mathbb\{E\}\_\{Z=f\(X\)\}\\left\[R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\-R\(Z\)\+\\lambda\\lVert Z\\rVert\_\{0\}\\right\],\(1\)where the global coding rate is defined as

R​\(Z\)=12​log​det\(Id\+dN​ϵ2​Z​Z⊤\)\.R\(Z\)=\\frac\{1\}\{2\}\\log\\det\\left\(I\_\{d\}\+\\frac\{d\}\{N\\epsilon^\{2\}\}ZZ^\{\\top\}\\right\)\.\(2\)By maximizing the log\-determinant of the representation covariance,R​\(Z\)R\(Z\)encourages the representations to occupy a large volume in the ambient space and thereby prevents global collapse\. Equivalently, under a fixed representation energy, maximizingR​\(Z\)R\(Z\)favors a balanced covariance spectrum and discourages the representations from concentrating along only a few directions\. Total coding\-rate maximization has accordingly been interpreted as a soft covariance regularization mechanism closely related to covariance\-based self\-supervised objectives such as VICReg\(Tonget al\.[2023](https://arxiv.org/html/2608.04213#bib.bib26)\)\.

SIGReg similarly encourages the embedding distribution to remain globally expanded by driving it toward a standard isotropic Gaussian\. In practice, SIGReg is evaluated on the embeddings used by the LeJEPA objective, which we denote byZ′Z^\{\\prime\}to distinguish them from the token matrixZZin Eq\. \([1](https://arxiv.org/html/2608.04213#Sx3.E1)\)\. The two variables are produced within the same representation\-learning pipeline, but need not denote the same finite\-sample matrix\. Our correspondence concerns the expansion effect that the two objectives impose on the learned representation distribution\.

SIGReg measures the discrepancy between random one\-dimensional projections of the embedding distribution and a standard Gaussian using an Epps–Pulley characteristic\-function statistic\(Balestriero and LeCun[2025](https://arxiv.org/html/2608.04213#bib.bib27)\)\. Around the isotropic Gaussian optimum, the discrepancy for centered embeddings can be decomposed as

ℒSIGReg​\(Z′\)=Cw​𝔼a∼Unif​\(𝕊d−1\)​\[\(a⊤​\(Σ−Id\)​a\)2\]\+ℒ\>2,\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}\(Z^\{\\prime\}\)=C\_\{w\}\\,\\mathbb\{E\}\_\{a\\sim\\mathrm\{Unif\}\(\\mathbb\{S\}^\{d\-1\}\)\}\\left\[\\bigl\(a^\{\\top\}\(\\Sigma\-I\_\{d\}\)a\\bigr\)^\{2\}\\right\]\+\\mathcal\{L\}\_\{\>2\},\(3\)whereCw\>0C\_\{w\}\>0,Σ\\Sigmadenotes the covariance ofZ′Z^\{\\prime\}, andℒ\>2\\mathcal\{L\}\_\{\>2\}collects deviations in skewness, kurtosis, and higher\-order statistics\. We denote the covariance\-level component in Eq\. \([3](https://arxiv.org/html/2608.04213#Sx3.E3)\) byℒSIGReg\(2\)\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\.

#### Theorem 1 \(Consistency of the expansion objectives\)\.

For centered embeddings with fixed normalized energytr⁡\(Σ\)=d\\operatorname\{tr\}\(\\Sigma\)=d, the coding\-rate objective and the covariance component of SIGReg share the same isotropic optimum:

arg⁡maxΣ⪰0tr⁡\(Σ\)=d⁡R​\(Σ\)=arg⁡minΣ⪰0tr⁡\(Σ\)=d⁡ℒSIGReg\(2\)​\(Σ\)=\{Id\}\.\\arg\\max\_\{\\begin\{subarray\}\{c\}\\Sigma\\succeq 0\\\\ \\operatorname\{tr\}\(\\Sigma\)=d\\end\{subarray\}\}R\(\\Sigma\)=\\arg\\min\_\{\\begin\{subarray\}\{c\}\\Sigma\\succeq 0\\\\ \\operatorname\{tr\}\(\\Sigma\)=d\\end\{subarray\}\}\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)=\\\{I\_\{d\}\\\}\.\(4\)Moreover, aroundΣ=Id\\Sigma=I\_\{d\}, the coding\-rate gapR​\(Id\)−R​\(Σ\)R\(I\_\{d\}\)\-R\(\\Sigma\)andℒSIGReg\(2\)​\(Σ\)\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)are both proportional, to leading order, to∥Σ−Id∥F2\\lVert\\Sigma\-I\_\{d\}\\rVert\_\{F\}^\{2\}\. The complete derivation is provided in Appendix[A](https://arxiv.org/html/2608.04213#A1)\.

Theorem 1 establishes that coding\-rate maximization and SIGReg implement a consistent expansion principle at the covariance level: both favor an isotropic representation covariance and penalize directional concentration around the optimum\. They are nevertheless not globally identical objectives\. The coding rate depends only on second\-order statistics, whereas SIGReg additionally constrains higher\-order deviations throughℒ\>2\\mathcal\{L\}\_\{\>2\}\. Thus, SIGReg preserves the covariance\-expansion effect ofR​\(Z\)R\(Z\)while strengthening it toward a full\-distribution isotropic Gaussian constraint\.

This correspondence allows the global expansion effect in Eq\. \([1](https://arxiv.org/html/2608.04213#Sx3.E1)\) to be incorporated into the LeJEPA training objective\. The compression and sparsity terms can then be retained explicitly for the subsequent derivation of the network architecture\.

### Optimizing the Compression and Sparsity Terms

With the global expansion term handled by the training loss, the remaining optimization problem is

minZ⁡Rc​\(Z∣U\[K\]\)\+λ​∥Z∥0\.\\min\_\{Z\}\\;R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\+\\lambda\\lVert Z\\rVert\_\{0\}\.\(5\)ForUk∈ℝd×pU\_\{k\}\\in\\mathbb\{R\}^\{d\\times p\}withUk⊤​Uk=IpU\_\{k\}^\{\\top\}U\_\{k\}=I\_\{p\}, the coding rate of the representations projected onto thekk\-th subspace is defined as

R​\(Uk⊤​Z\)=12​log​det\(Ip\+pN​ϵ2​Uk⊤​Z​Z⊤​Uk\),R\(U\_\{k\}^\{\\top\}Z\)=\\frac\{1\}\{2\}\\log\\det\\left\(I\_\{p\}\+\\frac\{p\}\{N\\epsilon^\{2\}\}U\_\{k\}^\{\\top\}ZZ^\{\\top\}U\_\{k\}\\right\),\(6\)and the corresponding compression rate is

Rc​\(Z∣U\[K\]\)=∑k=1KR​\(Uk⊤​Z\)\.R^\{c\}\(Z\\mid U\_\{\[K\]\}\)=\\sum\_\{k=1\}^\{K\}R\(U\_\{k\}^\{\\top\}Z\)\.\(7\)
To obtain a tractable sparsity surrogate, we replace theℓ0\\ell\_\{0\}penalty by anℓ1\\ell\_\{1\}penalty, as commonly done in sparse recovery\(Chenet al\.[1998](https://arxiv.org/html/2608.04213#bib.bib15)\), and introduce an auxiliary variableVV:

minZ,V⁡Rc​\(Z∣U\[K\]\)\+λ​∥V∥1,s\.t\.Z=V\.\\min\_\{Z,V\}\\;R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\+\\lambda\\lVert V\\rVert\_\{1\},\\quad\\mathrm\{s\.t\.\}\\;Z=V\.\(8\)Using the scaled\-form ADMM with dual variableWWand penalty parameterρ\>0\\rho\>0\(Boydet al\.[2011](https://arxiv.org/html/2608.04213#bib.bib29); Eckstein and Bertsekas[1992](https://arxiv.org/html/2608.04213#bib.bib19)\), the augmented Lagrangian, up to a term independent ofZZandVV, is

ℒρ​\(Z,V,W\)=Rc​\(Z∣U\[K\]\)\+λ​∥V∥1\+ρ2​∥Z−V\+W∥F2\.\\mathcal\{L\}\_\{\\rho\}\(Z,V,W\)=R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\+\\lambda\\lVert V\\rVert\_\{1\}\+\\frac\{\\rho\}\{2\}\\lVert Z\-V\+W\\rVert\_\{F\}^\{2\}\.\(9\)
The update ofZZ\-step is given by

Zt\+1=arg⁡minZ⁡Rc​\(Z∣U\[K\]\)\+ρ2​∥Z−Vt\+Wt∥F2\.Z^\{t\+1\}=\\arg\\min\_\{Z\}\\;R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\+\\frac\{\\rho\}\{2\}\\lVert Z\-V^\{t\}\+W^\{t\}\\rVert\_\{F\}^\{2\}\.\(10\)Following CRATE, we approximate this subproblem with a single linearized gradient step\(Yuet al\.[2023a](https://arxiv.org/html/2608.04213#bib.bib21); Ouyanget al\.[2015](https://arxiv.org/html/2608.04213#bib.bib18)\)\. In particular, the gradient of the compression rate is approximated as

∇ZRc​\(Z∣U\[K\]\)≈γ​\[Z−MSSA⁡\(Z∣U\[K\]\)\],\\nabla\_\{Z\}R^\{c\}\(Z\\mid U\_\{\[K\]\}\)\\approx\\gamma\\left\[Z\-\\operatorname\{MSSA\}\(Z\\mid U\_\{\[K\]\}\)\\right\],\(11\)whereγ\>0\\gamma\>0is determined by the coding precision, and multi\-head subspace self\-attention is defined as

MSSA⁡\(Z∣U\[K\]\)=∑k=1KUk​Uk⊤​Z​softmax⁡\(Z⊤​Uk​Uk⊤​Z\),\\operatorname\{MSSA\}\(Z\\mid U\_\{\[K\]\}\)=\\sum\_\{k=1\}^\{K\}U\_\{k\}U\_\{k\}^\{\\top\}Z\\,\\operatorname\{softmax\}\\left\(Z^\{\\top\}U\_\{k\}U\_\{k\}^\{\\top\}Z\\right\),\(12\)
Combining Eq\. \([11](https://arxiv.org/html/2608.04213#Sx3.E11)\) with the quadratic penalty, the gradient of theZZsubproblem is approximated by

γ​\[Z−MSSA⁡\(Z∣U\[K\]\)\]\+ρ​\(Z−Vt\+Wt\)\.\\gamma\\left\[Z\-\\operatorname\{MSSA\}\(Z\\mid U\_\{\[K\]\}\)\\right\]\+\\rho\\left\(Z\-V^\{t\}\+W^\{t\}\\right\)\.\(13\)A gradient step with step sizeη\>0\\eta\>0therefore yields

Zt\+1=\\displaystyle Z^\{t\+1\}=\{\}\(1−η​γ−η​ρ\)​Zt\+η​γ​MSSA⁡\(Zt∣U\[K\]t\)\\displaystyle\(1\-\\eta\\gamma\-\\eta\\rho\)Z^\{t\}\+\\eta\\gamma\\operatorname\{MSSA\}\(Z^\{t\}\\mid U\_\{\[K\]\}^\{t\}\)\+η​ρ​\(Vt−Wt\)\.\\displaystyle\+\\eta\\rho\(V^\{t\}\-W^\{t\}\)\.\(14\)
The update ofVV\-step is given by

Vt\+1=arg⁡minV⁡λ​∥V∥1\+ρ2​∥Zt\+1−V\+Wt∥F2,V^\{t\+1\}=\\arg\\min\_\{V\}\\;\\lambda\\lVert V\\rVert\_\{1\}\+\\frac\{\\rho\}\{2\}\\lVert Z^\{t\+1\}\-V\+W^\{t\}\\rVert\_\{F\}^\{2\},\(15\)whose solution is the proximal operator of theℓ1\\ell\_\{1\}norm\(Daubechieset al\.[2004](https://arxiv.org/html/2608.04213#bib.bib17); Beck and Teboulle[2009](https://arxiv.org/html/2608.04213#bib.bib14); Parikh and Boyd[2014](https://arxiv.org/html/2608.04213#bib.bib16)\),

Vt\+1=shrinkλ/ρ⁡\(Zt\+1\+Wt\)\.V^\{t\+1\}=\\operatorname\{shrink\}\_\{\\lambda/\\rho\}\(Z^\{t\+1\}\+W^\{t\}\)\.\(16\)
Following CRATE, we further impose a non\-negativity constraint on the intermediate representations\. Under this constraint, the proximal mapping reduces to the one\-sided soft\-thresholding operator, yielding

Vt\+1=ReLU⁡\(Zt\+1\+Wt−λ/ρ\)\.V^\{t\+1\}=\\operatorname\{ReLU\}\(Z^\{t\+1\}\+W^\{t\}\-\\lambda/\\rho\)\.\(17\)
This update preserves the sparse proximal interpretation while remaining compatible with the nonlinear activation commonly adopted in transformer architectures\.

### Attention\-Only White\-Box Transformer

![Refer to caption](https://arxiv.org/html/2608.04213v1/x2.png)Figure 2:One AoT\-ADMM layer as an unrolled ADMM iteration\. MSSA realizes the compression update, the ReLU proximal operator implements nonnegative shrinkage, and the scaled multiplier state tracks the consistency betweenZZandVV\. The three branch coefficients are initialized from the ADMM\-derived update and adapt dynamically during training\.Following the algorithm\-unrolling paradigm exemplified by learned sparse\-coding iterations\(Gregor and LeCun[2010](https://arxiv.org/html/2608.04213#bib.bib13)\), mapping the ADMM iteration indexttto the layer indexℓ\\ellconverts the iterative updates into the forward computation of one network layer:

\{Zℓ\+1=\(1−η​γ−η​ρ\)​Zℓ\+η​γ​MSSA⁡\(Zℓ∣U\[K\]ℓ\)\+η​ρ​\(Vℓ−Wℓ\),Vℓ\+1=ReLU⁡\(Zℓ\+1\+Wℓ−λ/ρ\),Wℓ\+1=Wℓ\+Zℓ\+1−Vℓ\+1\.\\left\\\{\\begin\{aligned\} Z^\{\\ell\+1\}=\{\}&\(1\-\\eta\\gamma\-\\eta\\rho\)Z^\{\\ell\}\+\\eta\\gamma\\operatorname\{MSSA\}\(Z^\{\\ell\}\\mid U\_\{\[K\]\}^\{\\ell\}\)\\\\ &\+\\eta\\rho\(V^\{\\ell\}\-W^\{\\ell\}\),\\\\\[2\.0pt\] V^\{\\ell\+1\}=\{\}&\\operatorname\{ReLU\}\\left\(Z^\{\\ell\+1\}\+W^\{\\ell\}\-\\lambda/\\rho\\right\),\\\\\[2\.0pt\] W^\{\\ell\+1\}=\{\}&W^\{\\ell\}\+Z^\{\\ell\+1\}\-V^\{\\ell\+1\}\.\\end\{aligned\}\\right\.\(18\)
The resulting network propagates the state triplet\(Zℓ,Vℓ,Wℓ\)\(Z^\{\\ell\},V^\{\\ell\},W^\{\\ell\}\)through depth, with each layer corresponding to one unrolled ADMM iteration\. As illustrated in Figure[2](https://arxiv.org/html/2608.04213#Sx3.F2), theZZ\-step performs MSSA\-based compression with an ADMM consistency correction, theVV\-step applies nonnegative shrinkage through a ReLU proximal operator, and theWW\-step updates the dual state to promote agreement betweenZℓ\+1Z^\{\\ell\+1\}andVℓ\+1V^\{\\ell\+1\}\. All layers share the same ADMM\-derived update form, while using independently parameterized subspace basesU\[K\]ℓU\_\{\[K\]\}^\{\\ell\}to adapt the compression model to the representation distribution at each depth\. StackingLLsuch iterations yields the complete attention\-only encoder\.

In implementation, the three branch coefficients in Eq\. \([18](https://arxiv.org/html/2608.04213#Sx3.E18)\) are initialized from the ADMM\-derived update and subsequently learned under a nonnegative unit\-sum constraint\. Figure[2](https://arxiv.org/html/2608.04213#Sx3.F2)therefore summarizes both the optimization interpretation and the forward computation of a single AoT\-ADMM layer\.

We embed the resulting encoder in LeJEPA\. Multi\-crop augmentations\(Caronet al\.[2021](https://arxiv.org/html/2608.04213#bib.bib5)\)generate global and local views of each image, which are mapped by the shared encoder and projection head to their corresponding embeddings\. A predictive consistency loss aligns the local and global embeddings with the mean global representation, while SIGReg regularizes the batch embedding distribution toward a standard isotropic Gaussian\. The total training objective is

ℒ=ℒpred\+α​ℒSIGReg,\\mathcal\{L\}=\\mathcal\{L\}\_\{\\text\{pred\}\}\+\\alpha\\,\\mathcal\{L\}\_\{\\text\{SIGReg\}\},\(19\)whereα\\alphabalances semantic invariance and distributional expansion\. The encoder, projection head, and layerwise subspace bases are optimized jointly end to end\.

## Experiments

### Experimental Setup

We conduct experiments on CIFAR\-10, CIFAR\-100\(Krizhevsky[2009](https://arxiv.org/html/2608.04213#bib.bib30)\), and ImageNet\-1K\(Denget al\.[2009](https://arxiv.org/html/2608.04213#bib.bib31)\)\. For in\-domain self\-supervised learning, CRATE and AoT\-ADMM are pretrained on CIFAR\-10 and CIFAR\-100 using the same LeJEPA objective and evaluated by linear probing with frozen backbones, following the standard self\-supervised evaluation protocol\(Chenet al\.[2020](https://arxiv.org/html/2608.04213#bib.bib3); Grillet al\.[2020](https://arxiv.org/html/2608.04213#bib.bib4)\)\. We report the top\-1 accuracy, parameters and flops as indicators of representation quality and model complexity\. We further pretrain AoT\-ADMM on ImageNet\-1K and fully fine\-tune the resulting backbone on CIFAR\-10 and CIFAR\-100 to evaluate large\-scale pretraining and transfer\. All experiments are conducted on a single NVIDIA A100 \(80GB\) GPU; detailed software and hardware configurations are provided in the supplementary material\.

For the direct comparison with existing attention\-only architectures, AoT\-ADMM and Wang et al\.’s AoT are evaluated at comparable model scales under the same LeJEPA training protocol\.

To further examine MLP redundancy beyond the white\-box CRATE architecture, we construct a standard ViT control by retaining its patch embedding, self\-attention, normalization layers, and classification head, while replacing the MLP in every Transformer block with a ReLU activation\. The resulting model is trained with knowledge distillation and compared with the original ViT in terms of Top\-1 accuracy and parameter count\.

For CIFAR pretraining, global and local views are32×3232\\times 32and16×1616\\times 16, the patch size is 8, and each image produces two global and six local views\. We train for 800 epochs with AdamW\(Loshchilov and Hutter[2019](https://arxiv.org/html/2608.04213#bib.bib32)\), a learning rate of5×10−45\\times 10^\{\-4\}, weight decay5×10−25\\times 10^\{\-2\}, batch size 256, and cosine decay\. For ImageNet\-1K, global and local views are128×128128\\times 128and64×6464\\times 64, patch size is 16, and each image produces two global and two local views\. We train for 200 epochs with learning rate5×10−45\\times 10^\{\-4\}and batch size 1024; the resulting backbone is then fully fine\-tuned on CIFAR\.

### Main Results

We first evaluate in\-domain linear probing against CRATE under exactly the same LeJEPA pretraining settings, with results summarized in Table[1](https://arxiv.org/html/2608.04213#Sx4.T1)\. Across Tiny, Small, and Base, AoT\-ADMM reduces the parameter count by approximately31%31\\%while maintaining competitive performance\. At Tiny scale, the accuracy is only0\.100\.10points lower on CIFAR\-100\. At Small scale, AoT\-ADMM improves CIFAR\-100 accuracy from61\.64%61\.64\\%to62\.41%62\.41\\%, a gain of0\.770\.77points\. At Base scale, the gaps narrow to0\.300\.30and0\.020\.02points on CIFAR\-10 and CIFAR\-100, respectively\.

Table 1:Linear\-probe Top\-1 accuracy, FLOPs, and parameters under LeJEPA self\-supervised pretraining\.Figure[3](https://arxiv.org/html/2608.04213#Sx4.F3)further compares the LeJEPA pretraining dynamics of CRATE and AoT\-ADMM\. Under identical training settings, the two models follow closely matched convergence trajectories and reach comparable final loss values after 800 epochs, indicating that removing the learnable ISTA dictionary does not adversely affect optimization or convergence\.

These results show that the ADMM\-derived attention\-only structure preserves the representation quality of CRATE while removing the learnable ISTA dictionary and substantially reducing parameters\. A consistent reduction in model parameters is observed across all scales: AoT\-ADMM removes 1\.77M \(3\.64M vs\. 5\.41M\), 3\.99M \(8\.11M vs\. 12\.10M\), and 7\.08M \(14\.36M vs\. 21\.44M\) parameters from Tiny, Small, and Base, respectively, corresponding to reductions of31\.1%31\.1\\%,31\.2%31\.2\\%, and31\.3%31\.3\\%\. Together with the near\-identical Base\-scale performance and closely matched pretraining dynamics, this indicates that the expansion function can be transferred to the self\-supervised loss without requiring the original ISTA dictionary in every layer\.

![Refer to caption](https://arxiv.org/html/2608.04213v1/x3.png)Figure 3:Comparison of LeJEPA pretraining loss between CRATE and AoT\-ADMM\. Under identical settings, the two models exhibit closely matched convergence trajectories and reach comparable final loss values after 800 epochs\.Table 2:ImageNet\-1K self\-supervised pretraining \(200 epochs\) followed by full CIFAR fine\-tuning\.We further assess the transferability of AoT\-ADMM under large\-scale self\-supervised pretraining by pretraining on ImageNet\-1K and subsequently fine\-tuning on CIFAR\-10 and CIFAR\-100\. As shown in Table[2](https://arxiv.org/html/2608.04213#Sx4.T2), AoT\-ADMM achieves95\.25%95\.25\\%accuracy on CIFAR\-10 and80\.26%80\.26\\%on CIFAR\-100 with 14\.84M parameters, compared with96\.01%96\.01\\%and81\.04%81\.04\\%for CRATE with 21\.92M parameters\. These results demonstrate that AoT\-ADMM maintains competitive transfer performance while using approximately32%32\\%fewer parameters, highlighting its improved parameter efficiency under large\-scale self\-supervised pretraining\.

### Comparison with AoT

Wang et al\.’s AoT derives an attention\-only architecture by interpreting representation learning as iterative denoising toward a mixture of low\-dimensional subspaces\(Wanget al\.[2025](https://arxiv.org/html/2608.04213#bib.bib28)\)\. In contrast, AoT\-ADMM is derived by combining LeJEPA\-based global expansion with ADMM optimization of the remaining compression and sparsity terms\. To isolate the effect of the architectural formulation, we compare the two models at nearly identical parameter counts under the same LeJEPA training protocol\.

As shown in Table[3](https://arxiv.org/html/2608.04213#Sx4.T3), AoT\-ADMM consistently outperforms Wang et al\.’s AoT across Tiny, Small, and Base configurations on both CIFAR\-10 and CIFAR\-100, with nearly identical parameter counts\. At the Tiny scale, our model achieves accuracy improvements over AoT of 4\.91%\\%\(84\.22%\\%vs\. 79\.31%\\%\) on CIFAR\-10 and 6\.55%\\%\(57\.38%\\%vs\. 50\.83%\\%\) on CIFAR\-100\. At the Small scale, the accuracy gains are 3\.87%\\%\(87\.23%\\%vs\. 83\.36%\\%\) on CIFAR\-10 and 5\.90%\\%\(62\.41%\\%vs\. 56\.51%\\%\) on CIFAR\-100\. At the Base scale, our model outperforms AoT by 4\.68%\\%\(88\.88%\\%vs\. 84\.20%\\%\) in accuracy on CIFAR\-10 and by 6\.06%\\%\(63\.54%\\%vs\. 57\.48%\\%\) on CIFAR\-100\. These results indicate that removing the MLP alone is not sufficient to obtain an effective attention\-only Transformer; the optimization principle underlying the architectural derivation also plays an important role\.

Table 3:Comparison with Wang et al\.’s AoT under the LeJEPA evaluation\.
### Black\-Box Transformer Validation

We further test whether the proximal update can remove MLP blocks from a standard Vision Transformer\(Dosovitskiyet al\.[2021](https://arxiv.org/html/2608.04213#bib.bib33)\)\. This control is not identical to the white\-box CRATE architecture: it retains the standard ViT patch embedding, self\-attention, normalization, and classifier, while replacing every MLP branch with a simple ReLU activation\. To compensate for the reduced capacity, we initialize from pretrained weights and use hard\-label knowledge distillation following DeiT\(Touvronet al\.[2021](https://arxiv.org/html/2608.04213#bib.bib34)\)\. A pretrained ViT\-Tiny teacher provides its highest\-confidence class as a pseudo\-label\. The student minimizes an equal\-weighted sum of cross\-entropy losses against the ground\-truth label and the teacher pseudo\-label\.

Table 4:CIFAR\-10 validation of standard ViTs and their MLP\-free variants using hard\-label knowledge distillation\.Table[4](https://arxiv.org/html/2608.04213#Sx4.T4)shows that removing all MLP blocks reduces the parameter count by approximately66%66\\%at both model scales\. For the Tiny model, the parameter count decreases from 5\.3M to 1\.8M, while CIFAR\-10 accuracy drops only slightly from92\.19%92\.19\\%to91\.37%91\.37\\%, corresponding to a gap of0\.820\.82percentage points\. For the Small model, the parameter count is reduced from 21\.4M to 7\.2M, yet the accuracy increases marginally from93\.58%93\.58\\%to93\.67%93\.67\\%\. Although these results do not establish that MLP blocks are universally redundant, they indicate that optimization\-derived proximal updates can substantially improve parameter efficiency while preserving competitive performance, providing a principled route toward MLP\-free black\-box Transformers\.

## Conclusion

We developed an attention\-only white\-box Transformer by jointly reconsidering self\-supervised learning and optimization\-based architecture derivation\. Our key theoretical observation is that LeJEPA’s SIGReg and the global expansion termR​\(Z\)R\(Z\)share a consistent optimum under the isotropic covariance constraint\. This connection allows the expansion term to be incorporated into the self\-supervised training loss, while the remaining compression and sparsity terms are used to derive the network architecture\. By applying ADMM and unrolling the resulting updates, we obtain an attention\-only architecture in which MSSA implements the compression step and ReLU realizes the sparsity update\. The resulting AoT\-ADMM requires neither the ISTA dictionary nor MLP layers, establishing a unified connection among the white\-box objective, the self\-supervised loss, and the forward architecture\.

Experiments under identical LeJEPA pretraining show that AoT\-ADMM maintains performance comparable to CRATE on CIFAR\-10 and CIFAR\-100 while reducing the parameter count by approximately31%31\\%\. After ImageNet\-1K self\-supervised pretraining and full CIFAR fine\-tuning, AoT\-ADMM achieves95\.25%95\.25\\%on CIFAR\-10 and80\.26%80\.26\\%on CIFAR\-100, compared with CRATE’s96\.01%96\.01\\%and81\.04%81\.04\\%, while reducing the parameter count from 21\.92M to 14\.84M\. At nearly identical parameter counts, AoT\-ADMM also consistently outperforms Wang et al\.’s AoT across different model scales, demonstrating that the optimization principle underlying an attention\-only architecture is important for its effectiveness\. Experiments on standard ViTs further show that the ReLU update can replace MLP blocks while preserving competitive performance, suggesting that the proposed formulation is not limited to white\-box architectures\. The theoretical analysis and empirical results both substantiate the partial redundancy of MLP modules in white\-box Transformer architectures under self\-supervised training, while providing further motivation for the development of attention\-only Transformer architectures\.

## References

- M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, and N\. Ballas \(2023\)Self\-supervised learning from images with a joint\-embedding predictive architecture\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 15619–15629\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p5.3)\.
- R\. Balestriero and Y\. LeCun \(2025\)LeJEPA: provable and scalable self\-supervised learning without the heuristics\.arXiv preprint arXiv:2511\.08544\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p5.3),[Optimizing the Global Expansion Term](https://arxiv.org/html/2608.04213#Sx3.SSx1.p3.6)\.
- A\. Beck and M\. Teboulle \(2009\)A fast iterative shrinkage\-thresholding algorithm for linear inverse problems\.SIAM Journal on Imaging Sciences2\(1\),pp\. 183–202\.External Links:[Document](https://dx.doi.org/10.1137/080716542)Cited by:[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p5.2)\.
- S\. Boyd, N\. Parikh, E\. Chu, B\. Peleato, and J\. Eckstein \(2011\)Distributed optimization and statistical learning via the alternating direction method of multipliers\.Foundations and Trends in Machine Learning3\(1\),pp\. 1–122\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p5.3),[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p2.7)\.
- M\. Caron, H\. Touvron, I\. Misra, H\. Jégou, J\. Mairal, P\. Bojanowski, and A\. Joulin \(2021\)Emerging properties in self\-supervised vision transformers\.InProceedings of the IEEE/CVF International Conference on Computer Vision,pp\. 9650–9660\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p1.1),[Attention\-Only White\-Box Transformer](https://arxiv.org/html/2608.04213#Sx3.SSx3.p4.2)\.
- S\. S\. Chen, D\. L\. Donoho, and M\. A\. Saunders \(1998\)Atomic decomposition by basis pursuit\.SIAM Journal on Scientific Computing20\(1\),pp\. 33–61\.External Links:[Document](https://dx.doi.org/10.1137/S1064827596304010)Cited by:[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p2.3)\.
- T\. Chen, S\. Kornblith, M\. Norouzi, and G\. Hinton \(2020\)A simple framework for contrastive learning of visual representations\.InProceedings of the 37th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.119,pp\. 1597–1607\.Cited by:[Experimental Setup](https://arxiv.org/html/2608.04213#Sx4.SSx1.p1.1)\.
- X\. Chen, S\. Xie, and K\. He \(2021\)An empirical study of training self\-supervised vision transformers\.InProceedings of the IEEE/CVF International Conference on Computer Vision,pp\. 9640–9649\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p1.1)\.
- I\. Daubechies, M\. Defrise, and C\. De Mol \(2004\)An iterative thresholding algorithm for linear inverse problems with a sparsity constraint\.Communications on Pure and Applied Mathematics57\(11\),pp\. 1413–1457\.External Links:[Document](https://dx.doi.org/10.1002/cpa.20042)Cited by:[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p5.2)\.
- J\. Deng, W\. Dong, R\. Socher, L\. Li, K\. Li, and L\. Fei\-Fei \(2009\)ImageNet: a large\-scale hierarchical image database\.InIEEE Conference on Computer Vision and Pattern Recognition,pp\. 248–255\.Cited by:[Experimental Setup](https://arxiv.org/html/2608.04213#Sx4.SSx1.p1.1)\.
- Y\. Dong, J\. Cordonnier, and A\. Loukas \(2021\)Attention is not all you need: pure attention loses rank doubly exponentially with depth\.InProceedings of the 38th International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.139,pp\. 2793–2803\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p1.1),[White\-Box Self\-Supervised Learning](https://arxiv.org/html/2608.04213#Sx2.SSx1.p1.2),[Attention\-Only Transformers](https://arxiv.org/html/2608.04213#Sx2.SSx2.p1.1)\.
- A\. Dosovitskiy, L\. Beyer, A\. Kolesnikov, D\. Weissenborn, X\. Zhai, T\. Unterthiner, M\. Dehghani, M\. Minderer, G\. Heigold, S\. Gelly, J\. Uszkoreit, and N\. Houlsby \(2021\)An image is worth 16x16 words: transformers for image recognition at scale\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p1.1),[Black\-Box Transformer Validation](https://arxiv.org/html/2608.04213#Sx4.SSx4.p1.1)\.
- J\. Eckstein and D\. P\. Bertsekas \(1992\)On the douglas–rachford splitting method and the proximal point algorithm for maximal monotone operators\.Mathematical Programming55\(1–3\),pp\. 293–318\.External Links:[Document](https://dx.doi.org/10.1007/BF01581204)Cited by:[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p2.7)\.
- K\. Gregor and Y\. LeCun \(2010\)Learning fast approximations of sparse coding\.InProceedings of the 27th International Conference on Machine Learning,pp\. 399–406\.Cited by:[Attention\-Only White\-Box Transformer](https://arxiv.org/html/2608.04213#Sx3.SSx3.p1.2)\.
- J\. Grill, F\. Strub, F\. Altché, C\. Tallec, P\. H\. Richemond, E\. Buchatskaya, C\. Doersch, B\. Avila Pires, Z\. D\. Guo, M\. G\. Azar, B\. Piot, K\. Kavukcuoglu, R\. Munos, and M\. Valko \(2020\)Bootstrap your own latent: a new approach to self\-supervised learning\.InAdvances in Neural Information Processing Systems,Vol\.33,pp\. 21271–21284\.Cited by:[Experimental Setup](https://arxiv.org/html/2608.04213#Sx4.SSx1.p1.1)\.
- B\. He and T\. Hofmann \(2024\)Simplifying transformer blocks\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p1.1),[White\-Box Self\-Supervised Learning](https://arxiv.org/html/2608.04213#Sx2.SSx1.p1.2),[Attention\-Only Transformers](https://arxiv.org/html/2608.04213#Sx2.SSx2.p1.1)\.
- K\. He, X\. Chen, S\. Xie, Y\. Li, P\. Dollár, and R\. Girshick \(2022\)Masked autoencoders are scalable vision learners\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 16000–16009\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p4.1)\.
- A\. Krizhevsky \(2009\)Learning multiple layers of features from tiny images\.Technical reportUniversity of Toronto\.Cited by:[Experimental Setup](https://arxiv.org/html/2608.04213#Sx4.SSx1.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.InInternational Conference on Learning Representations,Cited by:[Experimental Setup](https://arxiv.org/html/2608.04213#Sx4.SSx1.p4.7)\.
- Y\. Ouyang, Y\. Chen, G\. Lan, and Jr\. Pasiliao \(2015\)An accelerated linearized alternating direction method of multipliers\.SIAM Journal on Imaging Sciences8\(1\),pp\. 644–681\.External Links:[Document](https://dx.doi.org/10.1137/14095697X)Cited by:[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p3.3)\.
- D\. Pai, Z\. Wu, S\. Buchanan, Y\. Yu, and Y\. Ma \(2024\)Masked completion via structured diffusion with white\-box transformers\.InInternational Conference on Learning Representations,Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p3.1),[Introduction](https://arxiv.org/html/2608.04213#Sx1.p4.1),[White\-Box Self\-Supervised Learning](https://arxiv.org/html/2608.04213#Sx2.SSx1.p1.2)\.
- N\. Parikh and S\. Boyd \(2014\)Proximal algorithms\.Foundations and Trends in Optimization1\(3\),pp\. 127–239\.External Links:[Document](https://dx.doi.org/10.1561/2400000003)Cited by:[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p5.2)\.
- S\. Tong, Y\. Chen, Y\. Ma, and Y\. LeCun \(2023\)EMP\-SSL: towards self\-supervised learning in one training epoch\.arXiv preprint arXiv:2304\.03977\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p4.1),[White\-Box Self\-Supervised Learning](https://arxiv.org/html/2608.04213#Sx2.SSx1.p1.2),[Optimizing the Global Expansion Term](https://arxiv.org/html/2608.04213#Sx3.SSx1.p1.5)\.
- H\. Touvron, M\. Cord, M\. Douze, F\. Massa, A\. Sablayrolles, and H\. Jégou \(2021\)Training data\-efficient image transformers and distillation through attention\.InInternational Conference on Machine Learning,pp\. 10347–10357\.Cited by:[Black\-Box Transformer Validation](https://arxiv.org/html/2608.04213#Sx4.SSx4.p1.1)\.
- A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. Polosukhin \(2017\)Attention is all you need\.InAdvances in Neural Information Processing Systems,Vol\.30\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p1.1)\.
- P\. Wang, Y\. Lu, Y\. Yu, D\. Pai, Q\. Qu, and Y\. Ma \(2025\)Attention\-only transformers via unrolled subspace denoising\.InProceedings of the 42nd International Conference on Machine Learning,Proceedings of Machine Learning Research, Vol\.267,pp\. 63840–63859\.Cited by:[Attention\-Only Transformers](https://arxiv.org/html/2608.04213#Sx2.SSx2.p1.1),[Comparison with AoT](https://arxiv.org/html/2608.04213#Sx4.SSx3.p1.1)\.
- Z\. Wu, T\. Ding, Y\. Lu, D\. Pai, J\. Zhang, W\. Wang, Y\. Yu, Y\. Ma, and B\. D\. Haeffele \(2024\)Token statistics transformer: linear\-time attention via variational rate reduction\.arXiv preprint arXiv:2412\.17810\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p3.1)\.
- J\. Yang, X\. Li, D\. Pai, Y\. Zhou, Y\. Ma, Y\. Yu, and C\. Xie \(2024\)Scaling white\-box transformers for vision\.arXiv preprint arXiv:2405\.20299\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p3.1)\.
- J\. Yang and Y\. Zhang \(2011\)Alternating direction algorithms forℓ1\\ell\_\{1\}\-problems in compressive sensing\.SIAM Journal on Scientific Computing33\(1\),pp\. 250–278\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p5.3)\.
- Y\. Yu, S\. Buchanan, D\. Pai, T\. Chu, Z\. Wu, S\. Tong, B\. D\. Haeffele, and Y\. Ma \(2023a\)White\-box transformers via sparse rate reduction\.Advances in Neural Information Processing Systems36\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p2.4),[White\-Box Self\-Supervised Learning](https://arxiv.org/html/2608.04213#Sx2.SSx1.p1.2),[Optimizing the Global Expansion Term](https://arxiv.org/html/2608.04213#Sx3.SSx1.p1.3),[Optimizing the Compression and Sparsity Terms](https://arxiv.org/html/2608.04213#Sx3.SSx2.p3.3)\.
- Y\. Yu, K\. H\. R\. Chan, C\. You, C\. Song, and Y\. Ma \(2020\)Learning diverse and discriminative representations via the principle of maximal coding rate reduction\.InAdvances in Neural Information Processing Systems,Vol\.33,pp\. 9422–9434\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p4.1),[Optimizing the Global Expansion Term](https://arxiv.org/html/2608.04213#Sx3.SSx1.p1.3)\.
- Y\. Yu, D\. Pai, S\. Buchanan, Z\. Wu, S\. Tong, B\. D\. Haeffele, and Y\. Ma \(2023b\)Emergence of segmentation with minimalistic white\-box transformers\.arXiv preprint arXiv:2308\.16271\.Cited by:[Introduction](https://arxiv.org/html/2608.04213#Sx1.p3.1)\.

## Appendix

## Appendix AProof of Theorem 1

This appendix provides the complete proof of Theorem 1 stated in the main text, together with a global extension of the correspondence that it establishes between the two expansion objectives\.

To avoid collision with the notation of the main text \(γ\\gamma: compression\-gradient coefficient;η\\eta: unrolled step size; andλ\\lambda: sparsity weight\), we define

c=dϵ2\>0,θ=c1\+c∈\(0,1\),A=Σ−Id\.c=\\frac\{d\}\{\\epsilon^\{2\}\}\>0,\\qquad\\theta=\\frac\{c\}\{1\+c\}\\in\(0,1\),\\qquad A=\\Sigma\-I\_\{d\}\.Letz∈ℝdz\\in\\mathbb\{R\}^\{d\}denote the population embedding on which SIGReg is evaluated, written asZ′Z^\{\\prime\}in the main text, with

μ=𝔼​\[z\]=0,Σ=𝔼​\[z​z⊤\]\.\\mu=\\mathbb\{E\}\[z\]=0,\\qquad\\Sigma=\\mathbb\{E\}\\bigl\[zz^\{\\top\}\\bigr\]\.Centering and the fixed normalized\-energy condition of Theorem 1 are imposed by the normalization preceding the regularizer\. Throughout this appendix the two objectives are compared on

𝒮=\{Σ∈𝕊\+d:tr⁡\(Σ\)=d\}\.\\mathcal\{S\}=\\bigl\\\{\\Sigma\\in\\mathbb\{S\}^\{d\}\_\{\+\}:\\operatorname\{tr\}\(\\Sigma\)=d\\bigr\\\}\.SinceA=Σ−IdA=\\Sigma\-I\_\{d\}, everyΣ∈𝒮\\Sigma\\in\\mathcal\{S\}satisfiestr⁡\(A\)=0\\operatorname\{tr\}\(A\)=0, a fact used repeatedly below\. On the unconstrained cone𝕊\+d\\mathbb\{S\}^\{d\}\_\{\+\}one hasR​\(α​Σ\)→∞R\(\\alpha\\Sigma\)\\to\\inftyasα→∞\\alpha\\to\\inftyfor everyΣ≠0\\Sigma\\neq 0, soRRattains no maximum there\. Both optima exist on𝒮\\mathcal\{S\}, by Eqs\. \([A\.13](https://arxiv.org/html/2608.04213#A1.E13)\) and \([A\.18](https://arxiv.org/html/2608.04213#A1.E18)\) below\.

The coding rate, expressed as a function of the population covariance, is

R​\(Σ\)=12​log​det\(Id\+c​Σ\),R\(\\Sigma\)=\\frac\{1\}\{2\}\\log\\det\(I\_\{d\}\+c\\Sigma\),\(A\.1\)which is the population form of the empirical objective: writingΣ^=1n​Z​Z⊤\\hat\{\\Sigma\}=\\frac\{1\}\{n\}ZZ^\{\\top\}for the sample covariance, one has12​log​det\(Id\+dn​ϵ2​Z​Z⊤\)=12​log​det\(Id\+c​Σ^\)\\frac\{1\}\{2\}\\log\\det\\bigl\(I\_\{d\}\+\\frac\{d\}\{n\\epsilon^\{2\}\}ZZ^\{\\top\}\\bigr\)=\\frac\{1\}\{2\}\\log\\det\(I\_\{d\}\+c\\hat\{\\Sigma\}\)with the same constantc=d/ϵ2c=d/\\epsilon^\{2\}\.

### A\.1Covariance Component of SIGReg

Fora∼Unif⁡\(𝕊d−1\)a\\sim\\operatorname\{Unif\}\(\\mathbb\{S\}^\{d\-1\}\), let

ua=a⊤​z,φa​\(t\)=𝔼​\[ei​t​ua\],sa=a⊤​\(Σ−Id\)​a,u\_\{a\}=a^\{\\top\}z,\\qquad\\varphi\_\{a\}\(t\)=\\mathbb\{E\}\[e^\{\\mathrm\{i\}tu\_\{a\}\}\],\\qquad s\_\{a\}=a^\{\\top\}\(\\Sigma\-I\_\{d\}\)a,where the projections are taken unstandardized; fora⊤​Σ​a\>0a^\{\\top\}\\Sigma a\>0one has1\+sa=a⊤​Σ​a1\+s\_\{a\}=a^\{\\top\}\\Sigma a, and a per\-direction rescaling ofuau\_\{a\}replacesφa\\varphi\_\{a\}by the characteristic function ofua/1\+sau\_\{a\}/\\sqrt\{1\+s\_\{a\}\}, in whichsas\_\{a\}does not appear\. The Epps–Pulley statistic used by SIGReg is

ℒSIGReg=𝔼a​∫ℝ\|φa​\(t\)−e−t2/2\|2​w​\(t\)​𝑑t,\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}=\\mathbb\{E\}\_\{a\}\\int\_\{\\mathbb\{R\}\}\\left\|\\varphi\_\{a\}\(t\)\-e^\{\-t^\{2\}/2\}\\right\|^\{2\}w\(t\)\\,dt,\(A\.2\)with weight functionw≥0w\\geq 0positive Lebesgue\-almost everywhere and satisfying∫ℝ\(1\+t8\)​w​\(t\)​𝑑t<∞\\int\_\{\\mathbb\{R\}\}\(1\+t^\{8\}\)w\(t\)\\,dt<\\infty\. The moment condition givesCw<∞C\_\{w\}<\\inftyin Eq\. \([A\.6](https://arxiv.org/html/2608.04213#A1.E6)\); positivity almost everywhere is used in the first part of the proof\.

Letκm​\(a\)\\kappa\_\{m\}\(a\)denote themm\-th cumulant ofuau\_\{a\}and set

Δ\>2=supasupm≥3\|κm​\(a\)\|,Δ=max⁡\{supa\|sa\|,Δ\>2\}\.\\Delta\_\{\{\>\}2\}=\\sup\_\{a\}\\sup\_\{m\\geq 3\}\\lvert\\kappa\_\{m\}\(a\)\\rvert,\\qquad\\Delta=\\max\\bigl\\\{\\textstyle\\sup\_\{a\}\\lvert s\_\{a\}\\rvert,\\ \\Delta\_\{\{\>\}2\}\\bigr\\\}\.Around the standard isotropic Gaussian target, and assumingΔ\>2<∞\\Delta\_\{\{\>\}2\}<\\inftyso that the series may be integrated againstwwterm by term, the projected characteristic function admits the local expansion

φa​\(t\)−e−t2/2\\displaystyle\\varphi\_\{a\}\(t\)\-e^\{\-t^\{2\}/2\}\(A\.3\)=e−t2/2​\[−sa​t22\+∑m≥3κm​\(a\)​\(i​t\)mm\!\+O​\(Δ2\)\]\.\\displaystyle=e^\{\-t^\{2\}/2\}\\left\[\-\\frac\{s\_\{a\}t^\{2\}\}\{2\}\+\\sum\_\{m\\geq 3\}\\kappa\_\{m\}\(a\)\\frac\{\(\\mathrm\{i\}t\)^\{m\}\}\{m\!\}\+O\(\\Delta^\{2\}\)\\right\]\.
Squaring Eq\. \([A\.3](https://arxiv.org/html/2608.04213#A1.E3)\), the terms free of any cumulant of order three or higher reduce to

e−t2​sa2​t44,e^\{\-t^\{2\}\}\\frac\{s\_\{a\}^\{2\}t^\{4\}\}\{4\},\(A\.4\)while every remaining term carries at least one factorκm​\(a\)\\kappa\_\{m\}\(a\)withm≥3m\\geq 3\. We take this cumulant order as the grouping criterion, collecting the former intoℒSIGReg\(2\)\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}and the latter into the termℒ\>2\\mathcal\{L\}\_\{\>2\}used in the main text\. The grouping is not an orthogonal decomposition: the cross termssa​κm​\(a\)s\_\{a\}\\kappa\_\{m\}\(a\)enterℒ\>2\\mathcal\{L\}\_\{\>2\}at the same order inΔ\\Deltaassa2s\_\{a\}^\{2\}entersℒSIGReg\(2\)\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\. On the Gaussian slice

𝒢=\{𝒩​\(0,Σ\):Σ∈𝒮\},κm​\(a\)≡0\(m≥3\),\\mathcal\{G\}=\\bigl\\\{\\mathcal\{N\}\(0,\\Sigma\):\\Sigma\\in\\mathcal\{S\}\\bigr\\\},\\qquad\\kappa\_\{m\}\(a\)\\equiv 0\\quad\(m\\geq 3\),\(A\.5\)one hasℒ\>2=0\\mathcal\{L\}\_\{\>2\}=0and the two components separate exactly; elsewhere they separate up toO​\(Δ\>2\)O\(\\Delta\_\{\{\>\}2\}\)\. The qualifier “at the covariance level” in Theorem 1 refers to this restriction\.

Integrating Eq\. \([A\.4](https://arxiv.org/html/2608.04213#A1.E4)\) againstw​\(t\)w\(t\)and averaging overaagives

ℒSIGReg\(2\)​\(Σ\)=Cw​𝔼a∼Unif⁡\(𝕊d−1\)​\[\(a⊤​A​a\)2\],\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)=C\_\{w\}\\,\\mathbb\{E\}\_\{a\\sim\\operatorname\{Unif\}\(\\mathbb\{S\}^\{d\-1\}\)\}\\left\[\(a^\{\\top\}Aa\)^\{2\}\\right\],\(A\.6\)Cw=14​∫ℝt4​e−t2​w​\(t\)​𝑑t∈\(0,∞\)\.\\displaystyle C\_\{w\}=\\frac\{1\}\{4\}\\int\_\{\\mathbb\{R\}\}t^\{4\}e^\{\-t^\{2\}\}w\(t\)\\,dt\\in\(0,\\infty\)\.

### A\.2Proof of the Shared Optimum and Local Correspondence

###### Proof of Theorem 1\.

The integrand of Eq\. \([A\.2](https://arxiv.org/html/2608.04213#A1.E2)\) is nonnegative, soℒSIGReg=0\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}=0forcesφa​\(t\)=e−t2/2\\varphi\_\{a\}\(t\)=e^\{\-t^\{2\}/2\}for almost everyttand almost everya∈𝕊d−1a\\in\\mathbb\{S\}^\{d\-1\}\. Sincew\>0w\>0almost everywhere and both sides are continuous intt, the identity holds for everytt; continuity ofa↦φa​\(t\)a\\mapsto\\varphi\_\{a\}\(t\)extends it to everyaa\. Every one\-dimensional projection ofzzis then standard normal, and the Cramér–Wold device gives

ℒSIGReg=0⇔z∼𝒩​\(0,Id\)⟹Σ=Id,\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}=0\\iff z\\sim\\mathcal\{N\}\(0,I\_\{d\}\)\\implies\\Sigma=I\_\{d\},\(A\.7\)the minimum being attained on𝒮\\mathcal\{S\}because𝒩​\(0,Id\)\\mathcal\{N\}\(0,I\_\{d\}\)satisfiestr⁡\(Σ\)=d\\operatorname\{tr\}\(\\Sigma\)=d\. Equation \([A\.7](https://arxiv.org/html/2608.04213#A1.E7)\) characterizes the minimizer as a distribution;RRdepends onΣ\\Sigmaalone\.

Restricted to the covariance level, the same conclusion follows in closed form\. Fora∼Unif⁡\(𝕊d−1\)a\\sim\\operatorname\{Unif\}\(\\mathbb\{S\}^\{d\-1\}\), the fourth\-order spherical moments satisfy

𝔼​\[ai​aj​ak​al\]=δi​j​δk​l\+δi​k​δj​l\+δi​l​δj​kd​\(d\+2\)\.\\mathbb\{E\}\[a\_\{i\}a\_\{j\}a\_\{k\}a\_\{l\}\]=\\frac\{\\delta\_\{ij\}\\delta\_\{kl\}\+\\delta\_\{ik\}\\delta\_\{jl\}\+\\delta\_\{il\}\\delta\_\{jk\}\}\{d\(d\+2\)\}\.\(A\.8\)
Using Eq\. \([A\.8](https://arxiv.org/html/2608.04213#A1.E8)\),

𝔼a​\[\(a⊤​A​a\)2\]\\displaystyle\\mathbb\{E\}\_\{a\}\[\(a^\{\\top\}Aa\)^\{2\}\]=∑i,j,k,lAi​j​Ak​l​𝔼​\[ai​aj​ak​al\]\\displaystyle=\\sum\_\{i,j,k,l\}A\_\{ij\}A\_\{kl\}\\mathbb\{E\}\[a\_\{i\}a\_\{j\}a\_\{k\}a\_\{l\}\]=tr\(A\)2\+tr\(A⊤A\)\+tr\(A2\)d​\(d\+2\)\.\\displaystyle=\\frac\{\\operatorname\{tr\}\(A\)^\{2\}\+\\operatorname\{tr\}\(A^\{\\top\}A\)\+\\operatorname\{tr\}\(A^\{2\}\)\}\{d\(d\+2\)\}\.\(A\.9\)
SinceAAis symmetric,

tr⁡\(A⊤​A\)=tr⁡\(A2\)=∥A∥F2\.\\operatorname\{tr\}\(A^\{\\top\}A\)=\\operatorname\{tr\}\(A^\{2\}\)=\\lVert A\\rVert\_\{F\}^\{2\}\.\(A\.10\)
Together withtr⁡\(A\)=0\\operatorname\{tr\}\(A\)=0,

𝔼a​\[\(a⊤​A​a\)2\]=2d​\(d\+2\)​∥A∥F2\.\\mathbb\{E\}\_\{a\}\[\(a^\{\\top\}Aa\)^\{2\}\]=\\frac\{2\}\{d\(d\+2\)\}\\lVert A\\rVert\_\{F\}^\{2\}\.\(A\.11\)
Substituting Eq\. \([A\.11](https://arxiv.org/html/2608.04213#A1.E11)\) into Eq\. \([A\.6](https://arxiv.org/html/2608.04213#A1.E6)\),

ℒSIGReg\(2\)​\(Σ\)=2​Cwd​\(d\+2\)​∥Σ−Id∥F2\.\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)=\\frac\{2C\_\{w\}\}\{d\(d\+2\)\}\\lVert\\Sigma\-I\_\{d\}\\rVert\_\{F\}^\{2\}\.\(A\.12\)
SinceCw\>0C\_\{w\}\>0, the right\-hand side is nonnegative and vanishes if and only ifΣ=Id\\Sigma=I\_\{d\}, in agreement with Eq\. \([A\.7](https://arxiv.org/html/2608.04213#A1.E7)\)\. Hence

arg⁡minΣ∈𝒮⁡ℒSIGReg\(2\)​\(Σ\)=\{Id\}\.\\arg\\min\_\{\\Sigma\\in\\mathcal\{S\}\}\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)=\\\{I\_\{d\}\\\}\.\(A\.13\)
We next determine the maximizer of the coding rate\. Letν1,…,νd\\nu\_\{1\},\\ldots,\\nu\_\{d\}denote the eigenvalues ofΣ\\Sigma\. SinceΣ⪰0\\Sigma\\succeq 0andtr⁡\(Σ\)=d\\operatorname\{tr\}\(\\Sigma\)=d,

νi≥0,∑i=1dνi=d,\\nu\_\{i\}\\geq 0,\\qquad\\sum\_\{i=1\}^\{d\}\\nu\_\{i\}=d,\(A\.14\)and Eq\. \([A\.1](https://arxiv.org/html/2608.04213#A1.E1)\) becomes

R​\(Σ\)=12​∑i=1dlog⁡\(1\+c​νi\)\.R\(\\Sigma\)=\\frac\{1\}\{2\}\\sum\_\{i=1\}^\{d\}\\log\(1\+c\\nu\_\{i\}\)\.\(A\.15\)
The functiong​\(x\)=log⁡\(1\+c​x\)g\(x\)=\\log\(1\+cx\)is strictly concave on\[0,∞\)\[0,\\infty\)since

g′′​\(x\)=−c2\(1\+c​x\)2<0\.g^\{\\prime\\prime\}\(x\)=\-\\frac\{c^\{2\}\}\{\(1\+cx\)^\{2\}\}<0\.\(A\.16\)
Jensen’s inequality applied to Eq\. \([A\.15](https://arxiv.org/html/2608.04213#A1.E15)\) gives

R​\(Σ\)\\displaystyle R\(\\Sigma\)≤d2​log⁡\(1\+cd​∑i=1dνi\)\\displaystyle\\leq\\frac\{d\}\{2\}\\log\\left\(1\+\\frac\{c\}\{d\}\\sum\_\{i=1\}^\{d\}\\nu\_\{i\}\\right\)=d2​log⁡\(1\+c\)=R​\(Id\)\.\\displaystyle=\\frac\{d\}\{2\}\\log\(1\+c\)=R\(I\_\{d\}\)\.\(A\.17\)
By strict concavity ofgg, equality holds if and only ifν1=⋯=νd\\nu\_\{1\}=\\cdots=\\nu\_\{d\}, and the trace constraint forcesν1=⋯=νd=1\\nu\_\{1\}=\\cdots=\\nu\_\{d\}=1, i\.e\.Σ=Id\\Sigma=I\_\{d\}\. Hence

arg⁡maxΣ∈𝒮⁡R​\(Σ\)=\{Id\}\.\\arg\\max\_\{\\Sigma\\in\\mathcal\{S\}\}R\(\\Sigma\)=\\\{I\_\{d\}\\\}\.\(A\.18\)
Equations \([A\.13](https://arxiv.org/html/2608.04213#A1.E13)\) and \([A\.18](https://arxiv.org/html/2608.04213#A1.E18)\) give the shared\-optimum statement in Theorem 1\.

It remains to establish the local correspondence aroundΣ=Id\\Sigma=I\_\{d\}\. SubstitutingΣ=Id\+A\\Sigma=I\_\{d\}\+Ainto Eq\. \([A\.1](https://arxiv.org/html/2608.04213#A1.E1)\),

R​\(Id\+A\)\\displaystyle R\(I\_\{d\}\+A\)=12​log​det\(Id\+c​\(Id\+A\)\)\\displaystyle=\\frac\{1\}\{2\}\\log\\det\\left\(I\_\{d\}\+c\(I\_\{d\}\+A\)\\right\)=12​log​det\(\(1\+c\)​Id\+c​A\)\\displaystyle=\\frac\{1\}\{2\}\\log\\det\\left\(\(1\+c\)I\_\{d\}\+cA\\right\)=d2​log⁡\(1\+c\)\+12​log​det\(Id\+θ​A\)\.\\displaystyle=\\frac\{d\}\{2\}\\log\(1\+c\)\+\\frac\{1\}\{2\}\\log\\det\(I\_\{d\}\+\\theta A\)\.\(A\.19\)
On𝒮\\mathcal\{S\}one has∥A∥2≤d−1\\lVert A\\rVert\_\{2\}\\leq d\-1, so the condition∥A∥2<1/θ\\lVert A\\rVert\_\{2\}<1/\\thetadefines a proper subset of𝒮\\mathcal\{S\}containingIdI\_\{d\}\. On that subset the log\-determinant expansion gives

log​det\(Id\+θ​A\)=θ​tr⁡\(A\)−θ22​tr⁡\(A2\)\+O​\(∥A∥F3\)\.\\log\\det\(I\_\{d\}\+\\theta A\)=\\theta\\operatorname\{tr\}\(A\)\-\\frac\{\\theta^\{2\}\}\{2\}\\operatorname\{tr\}\(A^\{2\}\)\+O\\bigl\(\\lVert A\\rVert\_\{F\}^\{3\}\\bigr\)\.\(A\.20\)
Usingtr⁡\(A\)=0\\operatorname\{tr\}\(A\)=0andtr⁡\(A2\)=∥A∥F2\\operatorname\{tr\}\(A^\{2\}\)=\\lVert A\\rVert\_\{F\}^\{2\},

log​det\(Id\+θ​A\)=−θ22​∥A∥F2\+O​\(∥A∥F3\)\.\\log\\det\(I\_\{d\}\+\\theta A\)=\-\\frac\{\\theta^\{2\}\}\{2\}\\lVert A\\rVert\_\{F\}^\{2\}\+O\\bigl\(\\lVert A\\rVert\_\{F\}^\{3\}\\bigr\)\.\(A\.21\)
WithR​\(Id\)=d2​log⁡\(1\+c\)R\(I\_\{d\}\)=\\frac\{d\}\{2\}\\log\(1\+c\), Eqs\. \([A\.19](https://arxiv.org/html/2608.04213#A1.E19)\) and \([A\.21](https://arxiv.org/html/2608.04213#A1.E21)\) give

R​\(Id\)−R​\(Σ\)=θ24​∥Σ−Id∥F2\+O​\(∥Σ−Id∥F3\)\.R\(I\_\{d\}\)\-R\(\\Sigma\)=\\frac\{\\theta^\{2\}\}\{4\}\\lVert\\Sigma\-I\_\{d\}\\rVert\_\{F\}^\{2\}\+O\\left\(\\lVert\\Sigma\-I\_\{d\}\\rVert\_\{F\}^\{3\}\\right\)\.\(A\.22\)
Equation \([A\.12](https://arxiv.org/html/2608.04213#A1.E12)\) shows thatℒSIGReg\(2\)\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}is a positive multiple of∥Σ−Id∥F2\\lVert\\Sigma\-I\_\{d\}\\rVert\_\{F\}^\{2\}, and Eq\. \([A\.22](https://arxiv.org/html/2608.04213#A1.E22)\) shows that the leading term of the coding\-rate gap is likewise a positive multiple of it, so that

limΣ→IdℒSIGReg\(2\)​\(Σ\)R​\(Id\)−R​\(Σ\)=8​Cwθ2​d​\(d\+2\)∈\(0,∞\)\.\\lim\_\{\\Sigma\\to I\_\{d\}\}\\frac\{\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)\}\{R\(I\_\{d\}\)\-R\(\\Sigma\)\}=\\frac\{8C\_\{w\}\}\{\\theta^\{2\}d\(d\+2\)\}\\in\(0,\\infty\)\.\(A\.23\)
Equations \([A\.13](https://arxiv.org/html/2608.04213#A1.E13)\), \([A\.18](https://arxiv.org/html/2608.04213#A1.E18)\) and \([A\.23](https://arxiv.org/html/2608.04213#A1.E23)\) prove both statements of Theorem 1\. ∎

### A\.3Global Equivalence of the Two Expansion Terms

Theorem 1 gives a shared optimum and a leading\-order correspondence atΣ=Id\\Sigma=I\_\{d\}\. The following proposition extends the correspondence to all of𝒮\\mathcal\{S\}\.

###### Proposition 1\.

There exist constants0<k1≤k2<∞0<k\_\{1\}\\leq k\_\{2\}<\\infty, depending only ondd,ccandww, such that

k1​\(R​\(Id\)−R​\(Σ\)\)≤ℒSIGReg\(2\)​\(Σ\)≤k2​\(R​\(Id\)−R​\(Σ\)\)k\_\{1\}\\bigl\(R\(I\_\{d\}\)\-R\(\\Sigma\)\\bigr\)\\leq\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)\\leq k\_\{2\}\\bigl\(R\(I\_\{d\}\)\-R\(\\Sigma\)\\bigr\)\\qquad\(A\.24\)for allΣ∈𝒮\\Sigma\\in\\mathcal\{S\}\.

Consequently, for everyτ\>0\\tau\>0,

\{ℒSIGReg\(2\)≤k1​τ\}\\displaystyle\\bigl\\\{\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\\leq k\_\{1\}\\tau\\bigr\\\}\(A\.25\)⊆\{R​\(Id\)−R≤τ\}\\displaystyle\\subseteq\\bigl\\\{R\(I\_\{d\}\)\-R\\leq\\tau\\bigr\\\}⊆\{ℒSIGReg\(2\)≤k2​τ\}\.\\displaystyle\\subseteq\\bigl\\\{\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\\leq k\_\{2\}\\tau\\bigr\\\}\.
Moreover, let𝒫​\(X\)=X−tr⁡\(X\)d​Id\\mathcal\{P\}\(X\)=X\-\\frac\{\\operatorname\{tr\}\(X\)\}\{d\}I\_\{d\}denote the orthogonal projection onto the traceless symmetric matrices\. Then, withA=Σ−IdA=\\Sigma\-I\_\{d\}as in Eq\. \([A](https://arxiv.org/html/2608.04213#A1.Ex1)\),

𝒫​\[∇ΣℒSIGReg\(2\)\]=4​Cwd​\(d\+2\)​A,\\displaystyle\\mathcal\{P\}\\bigl\[\\nabla\_\{\\Sigma\}\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\\bigr\]=\\frac\{4C\_\{w\}\}\{d\(d\+2\)\}A,\(A\.26\)𝒫​\[∇Σ\(R​\(Id\)−R​\(Σ\)\)\]=θ22​A\+O​\(∥A∥F2\),\\displaystyle\\mathcal\{P\}\\bigl\[\\nabla\_\{\\Sigma\}\\bigl\(R\(I\_\{d\}\)\-R\(\\Sigma\)\\bigr\)\\bigr\]=\\frac\{\\theta^\{2\}\}\{2\}A\+O\\bigl\(\\lVert A\\rVert\_\{F\}^\{2\}\\bigr\),so that their inner product equals2​Cw​θ2d​\(d\+2\)​∥A∥F2\+O​\(∥A∥F3\)\\frac\{2C\_\{w\}\\theta^\{2\}\}\{d\(d\+2\)\}\\lVert A\\rVert\_\{F\}^\{2\}\+O\\bigl\(\\lVert A\\rVert\_\{F\}^\{3\}\\bigr\)\.

###### Proof\.

Define

ρ​\(Σ\)=ℒSIGReg\(2\)​\(Σ\)R​\(Id\)−R​\(Σ\),Σ∈𝒮∖\{Id\}\.\\rho\(\\Sigma\)=\\frac\{\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)\}\{R\(I\_\{d\}\)\-R\(\\Sigma\)\},\\qquad\\Sigma\\in\\mathcal\{S\}\\setminus\\\{I\_\{d\}\\\}\.
The eigenvalues of anyΣ∈𝒮\\Sigma\\in\\mathcal\{S\}lie in\[0,d\]\[0,d\], so𝒮\\mathcal\{S\}is compact\.

By Eqs\. \([A\.13](https://arxiv.org/html/2608.04213#A1.E13)\) and \([A\.18](https://arxiv.org/html/2608.04213#A1.E18)\), the numerator and denominator ofρ\\rhoare continuous and strictly positive on𝒮∖\{Id\}\\mathcal\{S\}\\setminus\\\{I\_\{d\}\\\}\. On the subset

∥A∥2≤1/2​θ,\\lVert A\\rVert\_\{2\}\\leq 1/2\\theta,the eigenvalues ofθ​A\\theta Aare bounded away from−1\-1, so the remainder in Eq\. \([A\.22](https://arxiv.org/html/2608.04213#A1.E22)\) is bounded byC​∥A∥F3C\\lVert A\\rVert\_\{F\}^\{3\}withCCdepending only onθ\\theta, hence uniformly in the directionA/∥A∥FA/\\lVert A\\rVert\_\{F\}\.

Dividing Eq\. \([A\.12](https://arxiv.org/html/2608.04213#A1.E12)\) by Eq\. \([A\.22](https://arxiv.org/html/2608.04213#A1.E22)\) therefore gives Eq\. \([A\.23](https://arxiv.org/html/2608.04213#A1.E23)\) as a limit along every sequenceΣ→Id\\Sigma\\to I\_\{d\}in𝒮\\mathcal\{S\}, andρ\\rhoextends continuously to𝒮\\mathcal\{S\}with value8​Cw/\(θ2​d​\(d\+2\)\)8C\_\{w\}/\\bigl\(\\theta^\{2\}d\(d\+2\)\\bigr\)atIdI\_\{d\}\.

A positive continuous function on a compact set attains a positive minimum and a finite maximum; takingk1=min𝒮⁡ρk\_\{1\}=\\min\_\{\\mathcal\{S\}\}\\rhoandk2=max𝒮⁡ρk\_\{2\}=\\max\_\{\\mathcal\{S\}\}\\rhogives Eq\. \([A\.24](https://arxiv.org/html/2608.04213#A1.E24)\)\. Equation \([A\.25](https://arxiv.org/html/2608.04213#A1.E25)\) follows from the two inequalities in Eq\. \([A\.24](https://arxiv.org/html/2608.04213#A1.E24)\) separately: ifℒSIGReg\(2\)​\(Σ\)≤k1​τ\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}\(\\Sigma\)\\leq k\_\{1\}\\tau, the left inequality givesk1​\(R​\(Id\)−R​\(Σ\)\)≤k1​τk\_\{1\}\\bigl\(R\(I\_\{d\}\)\-R\(\\Sigma\)\\bigr\)\\leq k\_\{1\}\\tauand hence the first inclusion; ifR​\(Id\)−R​\(Σ\)≤τR\(I\_\{d\}\)\-R\(\\Sigma\)\\leq\\tau, the right inequality gives the second\.

For the gradients, Eq\. \([A\.12](https://arxiv.org/html/2608.04213#A1.E12)\) gives∇ΣℒSIGReg\(2\)=4​Cwd​\(d\+2\)​A\\nabla\_\{\\Sigma\}\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}^\{\(2\)\}=\\frac\{4C\_\{w\}\}\{d\(d\+2\)\}A, which is already traceless bytr⁡\(A\)=0\\operatorname\{tr\}\(A\)=0\. From Eq\. \([A\.19](https://arxiv.org/html/2608.04213#A1.E19)\),∇Σ\[R​\(Id\)−R​\(Σ\)\]=−θ2​\(Id\+θ​A\)−1\\nabla\_\{\\Sigma\}\\bigl\[R\(I\_\{d\}\)\-R\(\\Sigma\)\\bigr\]=\-\\frac\{\\theta\}\{2\}\(I\_\{d\}\+\\theta A\)^\{\-1\}; expanding\(Id\+θ​A\)−1=Id−θ​A\+O​\(∥A∥F2\)\(I\_\{d\}\+\\theta A\)^\{\-1\}=I\_\{d\}\-\\theta A\+O\(\\lVert A\\rVert\_\{F\}^\{2\}\)and applying𝒫\\mathcal\{P\}, which annihilates the multiple ofIdI\_\{d\}, gives the second identity in Eq\. \([A\.26](https://arxiv.org/html/2608.04213#A1.E26)\)\. The stated inner product follows from⟨A,A⟩=∥A∥F2\\langle A,A\\rangle=\\lVert A\\rVert\_\{F\}^\{2\}\. ∎

Equation \([A\.24](https://arxiv.org/html/2608.04213#A1.E24)\) bounds each expansion objective by a fixed positive multiple of the other at every point of𝒮\\mathcal\{S\}, and Eq\. \([A\.25](https://arxiv.org/html/2608.04213#A1.E25)\) nests their sublevel sets: driving either quantity below a given threshold drives the other below a constant multiple of it, without restriction to a neighborhood ofIdI\_\{d\}\. Equation \([A\.26](https://arxiv.org/html/2608.04213#A1.E26)\) gives the corresponding first\-order statement: on the constraint tangent space the two projected gradients are positive multiples of the same matrix to leading order, so a descent step on either has positive inner product with a descent step on the other\. SIGReg can thus be used in place ofR​\(Z\)R\(Z\)as the expansion term in the LeJEPA objective\.

By Eq\. \([A\.7](https://arxiv.org/html/2608.04213#A1.E7)\), the full statistic retains the same optimum while characterizing it at the level of the entire distribution:RRis the covariance\-level projection of SIGReg\. Assigning the expansion role to SIGReg thus retains the constraint imposed byR​\(Z\)R\(Z\)and adds the higher\-order constraint carried byℒ\>2\\mathcal\{L\}\_\{\>2\}, which no function ofΣ\\Sigmaalone can express\. Equation \([A\.23](https://arxiv.org/html/2608.04213#A1.E23)\) fixes the scale: the proportionality constant decays asO​\(d−2\)O\(d^\{\-2\}\)while the coding\-rate gap carries the dimension\-free factorθ2/4\\theta^\{2\}/4, so converting the weight of the expansion term by this ratio preserves its effective strength\.

## Appendix BPseudocode of the Proposed Algorithm

Algorithm[1](https://arxiv.org/html/2608.04213#alg1)summarizes one end\-to\-end LeJEPA pretraining iteration with the proposed AoT\-ADMM encoder\. The three coefficients\(aℓ,bℓ,cℓ\)\(a\_\{\\ell\},b\_\{\\ell\},c\_\{\\ell\}\)of theZZ\-update are parameterized as a softmax over trainable logits, so that they remain positive and sum to one throughout training; the logits are initialized so that the softmax outputs\(1−η​γ−η​ρ,η​γ,η​ρ\)\(1\-\\eta\\gamma\-\\eta\\rho,\\eta\\gamma,\\eta\\rho\), recovering the unrolled ADMM step at initialization\. To stabilize the magnitudes of the recurrent states, bothZℓ\+1Z^\{\\ell\+1\}and the updated dual stateWℓ\+1W^\{\\ell\+1\}are subjected to feature\-wise RMS rescaling at each layer\. The thresholdτ=λ/ρ\\tau=\\lambda/\\rhoretains the proximal interpretation of the sparsity update\.

Algorithm 1LeJEPA Pretraining with the AoT\-ADMM EncoderInput: Mini\-batchℬ\\mathcal\{B\}; global and local augmentations; encoder depthLL Parameters: Subspace bases\{U\[K\]ℓ\}ℓ=0L−1\\\{U\_\{\[K\]\}^\{\\ell\}\\\}\_\{\\ell=0\}^\{L\-1\}, branch coefficients\{aℓ,bℓ,cℓ\}ℓ=0L−1\\\{a\_\{\\ell\},b\_\{\\ell\},c\_\{\\ell\}\\\}\_\{\\ell=0\}^\{L\-1\}, projection headgg, and SIGReg weightαsig\\alpha\_\{\\mathrm\{sig\}\} Output: Updated encoder and projection\-head parameters

1:Generate two global views and the prescribed local views for every image in

ℬ\\mathcal\{B\}\.

2:foreach augmented view

x\(v\)x^\{\(v\)\}do

3:

Z0,\(v\)←PatchEmbed⁡\(x\(v\)\)\+PosEmbedZ^\{0,\(v\)\}\\leftarrow\\operatorname\{PatchEmbed\}\(x^\{\(v\)\}\)\+\\operatorname\{PosEmbed\}
4:

V0,\(v\)←Z0,\(v\)V^\{0,\(v\)\}\\leftarrow Z^\{0,\(v\)\};

W0,\(v\)←0W^\{0,\(v\)\}\\leftarrow 0
5:for

ℓ=0\\ell=0to

L−1L\-1do

6:

Aℓ,\(v\)←MSSA⁡\(Zℓ,\(v\)∣U\[K\]ℓ\)A^\{\\ell,\(v\)\}\\leftarrow\\operatorname\{MSSA\}\(Z^\{\\ell,\(v\)\}\\mid U\_\{\[K\]\}^\{\\ell\}\)
7:

Zℓ\+1,\(v\)←aℓ​Zℓ,\(v\)\+bℓ​Aℓ,\(v\)\+cℓ​\(Vℓ,\(v\)−Wℓ,\(v\)\)Z^\{\\ell\+1,\(v\)\}\\leftarrow a\_\{\\ell\}Z^\{\\ell,\(v\)\}\+b\_\{\\ell\}A^\{\\ell,\(v\)\}\+c\_\{\\ell\}\(V^\{\\ell,\(v\)\}\-W^\{\\ell,\(v\)\}\)
8:

Vℓ\+1,\(v\)←ReLU⁡\(Zℓ\+1,\(v\)\+Wℓ,\(v\)−τ\)V^\{\\ell\+1,\(v\)\}\\leftarrow\\operatorname\{ReLU\}\(Z^\{\\ell\+1,\(v\)\}\+W^\{\\ell,\(v\)\}\-\\tau\)
9:

Wℓ\+1,\(v\)←Wℓ,\(v\)\+Zℓ\+1,\(v\)−Vℓ\+1,\(v\)W^\{\\ell\+1,\(v\)\}\\leftarrow W^\{\\ell,\(v\)\}\+Z^\{\\ell\+1,\(v\)\}\-V^\{\\ell\+1,\(v\)\}
10:endfor

11:

Y\(v\)←g​\(Pool⁡\(VL,\(v\)\)\)Y^\{\(v\)\}\\leftarrow g\\\!\\left\(\\operatorname\{Pool\}\(V^\{L,\(v\)\}\)\\right\)
12:endfor

13:Compute

ℒpred\\mathcal\{L\}\_\{\\mathrm\{pred\}\}from the local and global embeddings and their mean global target\.

14:Compute

ℒ←ℒpred\+αsig​ℒSIGReg\\mathcal\{L\}\\leftarrow\\mathcal\{L\}\_\{\\mathrm\{pred\}\}\+\\alpha\_\{\\mathrm\{sig\}\}\\mathcal\{L\}\_\{\\mathrm\{SIGReg\}\}\.

15:Update the encoder, layerwise subspace bases, branch coefficients, and projection head by backpropagation while maintaining the coefficient constraints\.

16:returnthe updated parameters\.

The inner loop is exactly one unrolled ADMM iteration per network layer: MSSA performs the compression update, ReLU implements the nonnegative soft\-thresholding step, and the dual stateWWtracks the consistency residual\. The outer part of the algorithm places this encoder inside the same multi\-crop LeJEPA objective used for all reported self\-supervised experiments\.