在深度学习中，人们一直在寻求更有效的方法以提高模型的性能和可解释性。近期，一群来自我国顶级高校和研究机构的研究人员提出了一种全新的Transformer结构——CRATE，该结构既具有出色的性能，又具备良好的可解释性。这一成果引发了人们对AI系统本质的深入思考：智能的本质是否就是压缩？CRATE的成功或许能为解决大模型的安全性问题提供一种可行的方案。来自UC伯克利、港大等机构的研究人员创新性地提出了一个“白盒”Transformer结构——CRATE。他们通过将数据从高维度分布压缩到低维结构分布，实现有效的表征，进而实现了具有竞争力的模型性能。这也引发了一个更为深远的讨论——智能的本质是否就是压缩？

AI领域的大佬们对大模型的安全性问题一直争论不休，原因在于神经网络的“黑箱”，使得所有人都难以理解其内在运作机制。实际上，如果能找到一种结合Transformer结构和功能性优势，同时具有良好可解释性的方法，那么大模型的安全性问题或许就能得到解决。为了降低未来超级AI对人类的潜在风险，Hinton等人大力倡导“AI末日论”，或许能引起公众的关注并推动达成共识。然而，要真正解决问题，必须从技术层面找到能够“彻底消除”AI风险的可行解决方案。因此，拆解大模型的“黑箱”就成为了关键的一步。

由马毅教授领导的来自加州大学伯克利分校（UC Berkeley）、清华大学产业创新与创业研究中心（TTIC）、上海科技大学、伊利诺伊大学香槟分校（UIUC）和约翰·霍普金斯大学（JHU）以及香港大学的 researchers 提出了一种名为「白盒」Transformer（CRATE）的创新方法。该方法可以在保持模型良好性能的同时，极大地提高模型的可解释性。

1. 压缩是AI系统的核心？

研究人员认为，要获得可解释的深度神经网络，必须从“第一性原理”出发，理解深度学习的本质。人工智能先驱 Hinton 在 20 世纪 90 年代提出了“深度学习的本质可能是压缩”的观点。许多 AI 大佬在各种场合对这个概念进行了经验性总结，并进一步扩展了这一理论。

例如，今年 8 月，OpenAI 首席科学家 Ilya Sutskever 在 UC Berkeley 的一个 AI 理论讲座上分享道：“压缩可能是学习的本质！”而马毅团队经过 5 年多的努力，完成了一篇长达 124 页的论文，更完整地阐述了这一理论。

更重要的是，他们基于这一理论设计出了可执行的算法，并在实践中取得了良好的性能表现。CRATE 在屏蔽任务上实现了具有竞争力的表现。

研究团队认为，数据表征学习的核心目标是从高维数据分布压缩到低维结构分布，以实现有效的表征。这种压缩可以通过“稀疏编码率减少”这个量化指标来衡量。研究团队利用简单的优化架构，将压缩和稀疏作为损失函数，可以迭代地将数据分布压缩到低维混合高斯分布模型，从而推导出类似于 Transformer 的神经网络结构。

这就是构建基于类Transformer的第一性原理。进一步证明压缩和去噪之间存在内在等价关系，就可以为构建Decoder提供理论依据，使编码器和解码器具有几乎相同的结构。研究团队的实验结果显示，尽管架构较为简单，CRATE在许多任务和数据集上都能与现有Transformer模型获得类似的表现，同时其每一层和操作都可以明确解释。分析结果表明，CRATE相对于标准Transformer确实具有更强的可解释性。

02 深度学习研究的新范式

而这个研究的另一方面意义在于，它直截了当地指出：“压缩就是一切。”在马毅教授看来，我们的研究表明：压缩似乎是当前人工智能系统的全部，包括GPT-4。剩下的问题是：仅压缩就能带来通用智能甚至意识吗？我敢打赌，答案显然是否定的。通过这项研究，所有人都对Transformer类型的AI系统获得了更加清晰的理解。而这就进一步说明，在外界看起来神奇和神秘的AI产品，只要背后的技术是基于Transformer，那么这些产品就不太可能超越纯机械数据压缩（编码）以及插值（解码）等功能。

之所以大众对于AI产品会有很多不切实际的幻想，可能根本原因就在于深度学习理论和实践长期脱节。而研究团队的这项工作，就是想弥合理论和实践之间的鸿沟，从而让AI产品背后的技术，理论都能严谨地结合在一起。

将理论层面的问题理清后，研究人员可以看到，现有的系统距离真正的智能系统还相差甚远，未来的提升空间仍然很大。现有的Transformer可能只是一种性价比不高且尚待改进的系统，后来者还需付出更多的努力！马毅教授在与我们的交流中提到：“当前沸沸扬扬的‘AI末日论’，直接催促他们紧急整合各种方法，利用有限的学术资源，力求全面、充分地对理论进行验证。” 如果在时间和资源允许的情况下，实验验证环节将更为充足，规模也可以更大。在完成这一工作并明确了现有方法的边界与本质后，研究团队将投身于更具挑战性和开创性的工作中。而“AI末日论”若最终导致人工智能研究受限或被扼杀，将违背所有人的利益；若由此引发可能的垄断，更是无法接受的。

03 白盒Transformer——CRATE

研究人员基于第一性原理构建了一个类似Transformer的架构，并将其命名为CRATE（Coding Rate Transformer）。在许多标准任务上，CRATE能够达到颇具竞争力的性能，同时具备诸多附加优势。CRATE是一种白盒（数学上可解释）Transformer架构，其中每一层都执行交替最小化算法的单个步骤，以优化稀疏率降低目标（sparse rate reduction objective）。具体而言，Norm功能旨在提高最终token表示的稀疏性。其函数定义如下：

其中，$F_{i}$是第$i$层的函数，$\theta_i$是第$i$层的参数，$\mu_{i}$是第$i$层的权重矩阵，$\sigma$是正弦激活函数，$\delta$是非线性激活函数。

层前向映射是一种方法，用于逐步优化稀疏率降低的目标。具体来说，这种映射将token表示为[Batch, Sequence Length, Dimension]。通过以下步骤，我们可以将其转换为多头子空间自注意力块和ISTA（迭代收缩阈值算法块，如图所示：

1. 将输入数据X编码为标记序列Z1后，CRATE构建了一个深度网络。该网络通过针对局部模型的连续压缩，将数据转换为低维子空间的规范配置分布，生成[Batch, Sequence Length, Dimension]的表示。

2. 针对全局字典进行稀疏化，生成[Batch, Sequence Length, Dimension - 1]的表示。

3. 重复堆叠这些块并通过反向传播训练模型参数，可以产生强大且可解释的数据表征。

4. 完整的架构仅由这些层的串联以及一些初始tokenizer和最终基于不同任务的架构组成（例如，头）。以下是CRATE的流程。它与常见的视觉Transformer原理相同。

5. 研究人员使用软最大交叉熵损失（soft-max cross entropy loss）来训练监督图像任务。以类似的缩放行为，他们使用经过训练的常用ViT获得了非常有竞争力的性能表现。例如，使用只有25%参数的ViT在ImageNet-1K上达到了80%以上的top-1准确率。

6. CRATE的一个有趣现象是，即使在监督方面进行训练，它也会学习对输入图像进行分割，并且这种分割可以通过注意力图轻松恢复，如图所示（类似于DINO）。

原始

这种分割以前只在DINO中使用复杂的自监督训练机制的类似Transformer的架构中看到，但在CRATE中，分割是监督训练的副产品。特别是，该模型在任何时候都不会获得任何先验分割信息。下面，研究人员展示了一些分割示例。CRATE的另一个显著的特性是注意力头自动携带语。这意味着CRATE的任何结果都能进行事后的解释。下面，研究人员将一些注意力头在几张图中的几种动物上的输出进行了可视化，显示了注意力头对应于动物的不同部分。而且结果表明，这种对应关系在不同动物的图片以及不同类别的动物图片中都是一致的。研究人员使用以下流程将CRATE扩展出了的能力。我们以扩散/最佳传输启发的方式构建解码器的每一层：如果我们认为以某种方式传输其输入分布的概率质量，那么被构造为该编码映射的近似逆。下面给出了完整的编码器和解码器层。CRATE架构的这种变体在屏蔽任务上实现了具有竞争力的性能，如下面的示例所示。此外，它还获得与经过训练的CRATE相同的涌现属性（如上所示）。研究人员通过对稀疏率降低的展开优化来推导编码器架构。优化稀疏率降低的表征ƒ是压缩和稀疏的，如下图所示，研究人员将它们描述为由编码器ƒ实现：在CRATE中，压缩运算符和稀疏化算子是稀疏率降低目标不同部分的近似（近端梯度步骤）。为了导出解码器架构，研究人员提出了一种新颖的结构化去噪扩散（structured diffusion）的方法。

在本文中，研究人员展示了一种新颖的结构化去噪扩散（structured diffusion）的方法，用于导出CRATE的解码器架构。CRATE架构是一种基于深度学习的图像分割模型，其显著特点是注意力头能够自动携带语言信息。研究人员发现，通过优化稀疏率降低的表征，可以有效压缩和解密图像特征，从而提高模型的性能。此外，研究人员还采用了一种扩散/最佳传输启发的方式构建解码器的每一层，使得CRATE的任何结果都能进行事后的解释。实验结果显示，CRATE架构的这种变体在屏蔽任务上具有竞争力，并且获得了与经过训练的CRATE相同的涌现属性。

Denoising-diffusion框架类似于广泛应用于图像数据生成模型的普通去噪扩散框架。其框架依赖于压缩算子和得分函数（例如在去噪扩散模型中所使用的定量连接，具体如下所示：编码器和解码器分别通过结构化去噪和扩散过程的离散化而得出）。值得注意的是，从展开优化导出的编码器和从结构化去噪导出的编码器具有相同的架构，如上述所提及的那样。下面是用于定义CRATE模型的示例代码。（参数为CRATE-Tiny指定）

```python

from model.crate import CRATE

dim = 384

n_heads = 6

depth = 12

model = CRATE(image_size=224, patch_size=16, num_classes=1000, dim=dim, depth=depth, heads=n_heads, dim_head=dim // n_heads)

```

要在ImageNet-1K上训练CRATE模型，请运行以下脚本（训练CRATE-tiny）：

```bash

python main.py --arch CRATE_tiny --batch-size 512 --epochs 200 --optimizer Lion --lr 0.0002 --weight-decay 0.05

```

这是在ImageNet-1K上训练CRATE-tiny的一个示例命令。

--print-freq 25

--data DATA_DIR, 并将DATA_DIR替换为[imagenet-folder with train and val folders].

python finetune.py --bs 256 --net CRATE_tiny --opt adamW --lr 5e-5 --n_epochs 200 --randomaug 1 --data cifar10 --ckpt_dir CKPT_DIR --data_dir DATA_DIR, 将CKPT_DIR替换为预训练CRATE权重的路径，并将DATA_DIR替换为CIFAR10数据集的路径。如果CKPT_DIR是None, 则此脚本用于在CIFAR10上通过随机初始化来训练CRATE。