arXiv:2608.04213cs.LG2026-08

用自监督学习构建纯注意力白盒Transformer,参数减少31%仍保持高精度。

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

论文配图:Attention-Only White-Box Transformer via LeJEPA-Based Self-Supervised Pretraining
图 1 · 摘自论文原文
  • 基于LeJEPA框架联合优化白盒Transformer,用注意力机制替代传统MLP结构。
  • 在CIFAR-10和CIFAR-100上准确率达88.88%和63.54%,仅损失0.3%性能。
  • 为标准ViT中MLP模块冗余提供新证据,适合关注模型压缩与架构简化者。

现有白盒网络自监督学习研究通常将优化算法与自监督范式解耦。本文从统一视角重新审视二者:基于LeJEPA的自监督框架假设最优嵌入分布为各向同性高斯分布,这在概念上等价于稀疏率减小目标中的展开项 $R(Z)$。据此,我们利用LeJEPA范式优化 $R(Z)$,并采用交替方向乘子法(ADMM)将剩余项 $R^{c}(Zackslashmid U_{[K]})+λackslashlVert Z Vert_{0}$ 转化为纯注意力结构,去除原设计中的ISTA结构或MLP层。实验表明,在基础规模下,本方法在CIFAR-10和CIFAR-100上的分类准确率分别为88.88%和63.54%,而原白盒模型CRATE分别为89.18%和63.56%,性能接近但参数量减少约31%。此外,我们进一步发现,在知识蒸馏下将标准ViT中所有MLP块替换为ReLU激活可减少约66%参数,同时保持竞争力,提示标准ViT中可能存在显著的MLP冗余。

原文摘要 · Abstract (English)

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.

Transformer自监督模型压缩注意力

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。