基于LeJEPA自监督预训练的纯注意力白盒Transformer

# 基于LeJEPA自监督预训练的纯注意力白盒Transformer  
Source: https://arxiv.org/html/2608.04213

###### 摘要  
现有的白盒网络自监督学习研究通常将白盒网络的推导（通过优化算法）与自监督学习范式解耦。在本工作中，我们转而从联合视角重新审视这两个组成部分。基于LeJEPA的自监督框架假设各向同性高斯分布是下游任务的最优嵌入分布，这在概念上等价于指导白盒Transformer优化的稀疏率降低目标中的扩展项R\(Z\)。基于这一观察，我们使用LeJEPA自监督范式优化R\(Z\)，并通过交替方向乘子法（ADMM）将余下的压缩项R^c\(Z∣U_[K]\)和稀疏项λ∥Z∥_0推导为一个纯注意力Transformer，从而省去原始设计中的ISTA结构或MLP层。实验结果表明，在LeJEPA自监督范式下，我们的纯注意力白盒Transformer在Base规模下于CIFAR-10上达到88.88%的分类准确率，在CIFAR-100上达到63.54%；而原始白盒Transformer CRATE在CIFAR-10上达到89.18%，在CIFAR-100上达到63.56%。我们的模型在取得竞争性性能的同时，参数量减少约31%。在白盒设置之外，我们进一步研究了标准ViT，发现在知识蒸馏下用ReLU激活替换所有MLP块可以移除约66%的参数并保持竞争性准确率，这促使我们进一步研究标准ViT架构中MLP模块的潜在冗余性。

## 引言  
Transformer已成为现代深度学习的基础架构，并在计算机视觉、自然语言处理等领域取得了显著性能（Vaswani等人, 2017 (https://arxiv.org/html/2608.04213#bib.bib1); Dosovitskiy等人, 2021 (https://arxiv.org/html/2608.04213#bib.bib33)）。然而，Transformer架构及其众多变体的设计在很大程度上仍是经验性的，缺乏严格的数学解释（He和Hofmann, 2024 (https://arxiv.org/html/2608.04213#bib.bib12); Dong等人, 2021 (https://arxiv.org/html/2608.04213#bib.bib11); Caron等人, 2021 (https://arxiv.org/html/2608.04213#bib.bib5); Chen等人, 2021 (https://arxiv.org/html/2608.04213#bib.bib10)）。这在很大程度上阻碍了具有更高效率或可解释性的新Transformer变体的开发。近年来，以编码率Transformer（CRATE）为代表的白盒Transformer架构为上述问题提供了一个有前景的新视角。CRATE以稀疏率降低原理为基础，通过扩展项R\(Z\)、压缩项R^c\(Z∣U_[K]\)和稀疏诱导项λ∥Z∥_0之间的协同作用来刻画理想表示（Yu等人, 2023a (https://arxiv.org/html/2608.04213#bib.bib21)）。在由此展开的网络中，压缩项通过单步梯度下降近似，从而产生多头子空间自注意力（MSSA）模块。同时，扩展项和稀疏项通过引入正交稀疏字典D，并结合迭代收缩阈值算法（ISTA）模块进行联合优化。这种展开框架将核心网络组件的计算明确解释为求解良定义目标子问题的迭代步骤，从而为Transformer提供了清晰统一的数学解释。在此基础上，白盒Transformer架构近年来持续演进。CRATE-α改进了稀疏编码模块和训练策略，显著增强了模型在大规模视觉任务上的可扩展性（Yang等人, 2024 (https://arxiv.org/html/2608.04213#bib.bib22)）。Token Statistics Transformer（ToST）则从最大编码率降低目标的变分公式中推导出具有线性和时间空间复杂度的可解释注意力机制（Wu等人, 2024 (https://arxiv.org/html/2608.04213#bib.bib23)）。与此同时，这一研究方向已扩展到无监督目标分割、语言建模和长序列建模等多种任务，凸显了白盒网络在可解释性、计算效率和跨领域适用性方面的广阔潜力（Yu等人, 2023b (https://arxiv.org/html/2608.04213#bib.bib24); Pai等人, 2024 (https://arxiv.org/html/2608.04213#bib.bib25)）。

白盒网络旨在直接从数据中揭示token表示背后的低维子空间结构。由于编码率目标无需语义标签即可构造，因此为自监督表示学习提供了一条可能路径（Yu等人, 2020 (https://arxiv.org/html/2608.04213#bib.bib2)）。近年来的几项研究尝试弥合这一差距。例如，CRATE-MAE将白盒编码器和解码器融入掩码自编码框架（Pai等人, 2024 (https://arxiv.org/html/2608.04213#bib.bib25); He等人, 2022 (https://arxiv.org/html/2608.04213#bib.bib8)），而EMP-SSL则将总编码率R\(Z\)引入自监督目标以促进表示扩展（Tong等人, 2023 (https://arxiv.org/html/2608.04213#bib.bib26)）。然而，这些方法仍存在明显局限：前者的重建损失缺乏清晰的理论解释，后者则保留了黑盒网络架构。这些局限引发了一个根本性问题：如何在自监督范式下训练模型，同时保持白盒架构与其学习目标之间的理论自洽？为解决这一问题，我们重新审视白盒目标的三个组成部分：扩展项被转移至LeJEPA自监督目标中进行优化，而压缩项和稀疏项则通过交替方向乘子法（ADMM）（Boyd等人, 2011 (https://arxiv.org/html/2608.04213#bib.bib29); Yang和Zhang, 2011 (https://arxiv.org/html/2608.04213#bib.bib35)）优化，并展开为纯注意力Transformer架构。我们的方法建立在白盒目标中的扩展项与LeJEPA施加的各向同性约束之间在表示空间中的紧密对应关系之上（Balestriero和LeCun, 2025 (https://arxiv.org/html/2608.04213#bib.bib27); Assran等人, 2023 (https://arxiv.org/html/2608.04213#bib.bib9)）。基于这一对应关系，我们用LeJEPA目标替换扩展项R\(Z\)，同时保留子空间压缩项R^c\(Z∣U_[K]\)和稀疏项λ∥Z∥_0，从而构造一个新的优化目标。我们随后使用ADMM对该目标进行优化。展开所得ADMM迭代，我们得到一个仅由注意力模块组成的白盒自监督架构。我们将这一由ADMM导出的纯注意力Transformer称为AoT-ADMM。AoT-ADMM的整体自监督训练框架如图1 (https://arxiv.org/html/2608.04213#Sx2.F1)所示。我们的主要贡献总结如下：

- 我们提出了一种基于LeJEPA自监督目标的纯注意力白盒Transformer。该模型使用LeJEPA自监督学习范式优化最大稀疏率降低中的扩展项R\(Z\)，而其余压缩项和稀疏项遵循白盒Transformer推导范式，从而将自监督学习与白盒优化目标结合起来。
- 我们使用LeJEPA自监督学习范式优化扩展项R\(Z\)，并使用ADMM优化其余白盒目标项。压缩项对应Transformer的注意力结构，而稀疏项通过ReLU算子得到闭式解，最终得到一个纯注意力白盒Transformer。
- 实验结果表明，AoT-ADMM在参数量减少约31%的情况下取得了与CRATE相当的性能。在Base规模下，我们的模型在CIFAR-10上达到88.88%，在CIFAR-100上达到63.54%，而CRATE分别为89.18%和63.56%。在相似参数量下，它也持续优于AoT。在标准ViT上，知识蒸馏下将所有MLP块替换为ReLU激活可移除约66%的参数并保持竞争性准确率，进一步表明MLP模块的潜在冗余性。

## 相关工作

### 白盒自监督学习  
传统Transformer模块通常通过经验性试错来选取（He和Hofmann, 2024 (https://arxiv.org/html/2608.04213#bib.bib12); Dong等人, 2021 (https://arxiv.org/html/2608.04213#bib.bib11)），架构与其学习目标之间缺乏严格的数学对应。CRATE则提供了一种替代方案：从稀疏率降低中推导其基本算子，使得每个网络操作都有明确的优化解释（Yu等人, 2023a (https://arxiv.org/html/2608.04213#bib.bib21)）。这一理论特性使白盒模型天然适合自监督表示学习。然而，以往尝试并未为训练目标和网络架构建立一致的理论来源。CRATE-MAE采用白盒编码器，但仍依赖不可解释的重建损失（Pai等人, 2024 (https://arxiv.org/html/2608.04213#bib.bib25)），而EMP-SSL引入了总编码率R\(Z\)，但继续采用黑盒架构（Tong等人, 2023 (https://arxiv.org/html/2608.04213#bib.bib26)）。相比之下，我们首先建立各向同性高斯先验与编码率扩展项R\(Z\)之间的理论联系，并将其作为训练目标，然后通过求解其余白盒目标项推导完整网络架构。据我们所知，这是第一个训练损失和前向计算均源于统一白盒理论框架的自监督公式。

### 纯注意力Transformer  
Transformer简化研究表明，某些块组件可以被移除或重新排列（He和Hofmann, 2024 (https://arxiv.org/html/2608.04213#bib.bib12)）。然而，朴素地移除前馈分支会削弱表示多样性，且理论分析已识别出纯注意力堆栈中的秩坍缩行为（Dong等人, 2021 (https://arxiv.org/html/2608.04213#bib.bib11)）。Wang等人提供了一项有趣的推导：他们将token建模为低秩高斯的噪声混合，将表示学习视为将噪声token压缩回其子空间，将多头子空间自注意力视作迭代去噪算子，并将该算子展开为仅由注意力和残差连接构成的纯注意力白盒Transformer（AoT）（Wang等人, 2025 (https://arxiv.org/html/2608.04213#bib.bib28)）。他们的目标是子空间去噪，因此纯注意力形式源于恢复干净表示，而非优化完整的表示学习目标。相比之下，我们的层来自不同的优化问题：其纯注意力架构和训练过程都自然地源于白盒目标。具体而言，我们保留稀疏率降低目标，将其扩展项移入训练损失，并通过展开压缩项和稀疏项解的迭代来获得纯注意力架构。

## 方法  
本节给出在基于LeJEPA的自监督预训练下推导纯注意力白盒Transformer的完整过程。如图1 (https://arxiv.org/html/2608.04213#Sx2.F1)所示，所提框架从稀疏率降低原理出发，同时推导训练目标和网络架构。我们首先建立LeJEPA中的草图式各向同性高斯正则化（SIGReg）与全局扩展项R\(Z\)之间的联系，从而使原始目标能够分离为损失层和架构层两个部分。然后，我们对余下的压缩项和稀疏项应用ADMM，得到一个三状态递归。最后，将这些ADMM迭代沿网络深度展开，构建纯注意力前向架构。

### 优化全局扩展项  
按照最大编码率降低及其稀疏白盒Transformer公式（Yu等人, 2020 (https://arxiv.org/html/2608.04213#bib.bib2), 2023a (https://arxiv.org/html/2608.04213#bib.bib21)），稀疏率降低将期望的表示刻画为全局扩展、在低维子空间内紧凑且稀疏。设Z=f\(X\)∈R^{d×N}表示N个token表示，U_[K]=\(U_k\)_{k=1}^K表示低维子空间基的集合。相应目标为  
\[
\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)  
其中全局编码率定义为  
\[
R(Z)=\frac{1}{2}\log\det\left(I_d+\frac{d}{N\epsilon^2}ZZ^\top\right).
\]
(2)  
通过最大化表示协方差的对数行列式，R\(Z\)鼓励表示在环境空间中占据较大体积，从而防止全局坍缩。等价地，在固定表示能量的条件下，最大化R\(Z\)有利于平衡的协方差谱，并阻止表示仅沿少数方向集中。因此，总编码率最大化被解释为一种软协方差正则化机制，与VICReg等基于协方差的自监督目标密切相关（Tong等人, 2023 (https://arxiv.org/html/2608.04213#bib.bib26)）。SIGReg同样通过将嵌入分布推向标准各向同性高斯来鼓励其保持全局扩展。在实践中，SIGReg在LeJEPA目标所使用的嵌入上求值，我们将其记为Z′，以区别于式(1)中的token矩阵Z。这两个变量在同一表示学习流程中产生，但不必表示同一个有限样本矩阵。我们的对应关系涉及两个目标对学习到的表示分布施加的扩展效应。SI