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

2026-08-04

AI Summary

This work addresses a key limitation in existing self-supervised learning approaches for white-box networks, which often decouple architectural derivation from learning objectives. The paper proposes the first unified optimization framework that jointly integrates LeJEPA-based self-supervised learning with white-box Transformer design. By solving a sparse rate-distortion objective via the Alternating Direction Method of Multipliers (ADMM), the method yields a pure attention architecture composed solely of attention mechanisms, achieving the first MLP-free white-box Vision Transformer (ViT). This finding reveals redundancy in the standard ViT’s MLP blocks. The resulting model attains competitive performance—88.88% and 63.54% accuracy on CIFAR-10 and CIFAR-100, respectively—while reducing parameter count by 31% relative to the baseline. Notably, even after removing the MLP entirely, yielding a 66% model compression, the architecture maintains strong performance.

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 term $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 optimize $R(Z)$, and derive the remaining terms $R^{c}(Z\mid U_{[K]})+λ\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 of $88.88\%$ on CIFAR-10 and $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 of $89.18\%$ on CIFAR-10 and $63.56\%$ on CIFAR-100. Our model achieves competitive performance while reducing the parameter count by roughly $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.

Problem

• white-box Transformer
• self-supervised learning
• attention-only architecture
• MLP redundancy
• parameter efficiency

Innovation

• LeJEPA
• white-box Transformer
• attention-only architecture
• self-supervised pretraining
• MLP redundancy