提出极简版新架构

编辑部 发自 凹非寺

量子位 | 公众号 QbitAI

CV大神何恺明，也来搞扩散模型（Diffusion Model）了！

大神最新论文刚刚挂上arXiv，还是热乎的：解构扩散模型，提出一个高度简化的新架构l-DAE（小写的L）。

并且通过与何恺明在视觉自监督学习领域的代表作MAE（Masked Autoencoder）对比，更好地理解了扩散模型内部的工作原理。

不仅如此，这篇论文还发现在扩散模型中，去噪过程比扩散过程更重要。

这项工作阵容非常豪华，不仅有何恺明坐镇，合著作者中还有纽约大学计算机科学助理教授、CV大牛谢赛宁。

以及曾和他共同发表ConvNeXT工作的刘壮——他同时是DenseNet的共同一作。

给扩散模型开刀

团队认为，尽管去噪扩散模型在生成任务上表现出色，但它们在表示学习方面的能力尚未得到充分探索。

为此，他们找到一个新颖的研究方法：

希望通过解构扩散模型，将其逐步转化为类似于MAE的架构，以更深入地理解DDM（去噪扩散模型）在自监督学习中的表示学习能力。

先来一图概括解构过程：

（如果你也觉得这个图很眼熟，没错，就是谢赛宁在代表作ConvNeXT中使用的同款。）

以使用了VQGAN tokenizer的DiT模型作为基线，接下来介绍团队如何一刀一刀把它改造成更简化版本的。

• 移除类别条件

首先移除扩散模型中基于类别标签的条件，很显然，有条件标签与自监督学习的目标不符。

这一步骤显著提高了线性探测（linear probing）的准确率。

• 解构VQGAN

接下来移除原本采用的VQGan感知损失和对抗损失，在自监督学习中，这两种损失函数都不是必须的。

感知损失涉及到对类别标签的依赖，这与自监督学习的目标不符；对抗损失涉及对生成过程的优化，而不是直接学习数据的表示。

移除两种损失函数后，模型表现受到影响下降，但下一步又一把拉了回来。

• 替换噪声调度器

在原始扩散模型中，噪声调度通常遵循一个复杂的时间步序列。团队采用了一种更简单的线性衰减策略，让噪声水平在训练过程中线性地从最大值衰减到零，而不是遵循复杂的非线性衰减路径。

这种简化的噪声调度策略使得模型在训练过程中更多地关注于较干净的数据，而不是在高度噪声的数据上花费太多计算资源。

• 简化tokenizer

作者进一步改造了扩散模型中的tokenizer，这是将输入图像映射到潜在空间的关键组件。

他们比较了几种不同的方法，包括卷积变分自编码器（conv. VAE）、基于块的变分自编码器（patch-wise VAE）、基于块的自编码器（patch-wise AE）和基于块的主成分分析（patch-wis PCA）。

最终发现，即使是简单的PCA也能有效地工作。

通过逆PCA（inverse PCA）将输入图像投影到潜在空间，添加噪声，然后再将噪声图像投影回图像空间。这种方法允许模型直接在图像上进行操作，而不需要tokenizer。

• 改变预测目标

与现代扩散模型通常预测噪声不同，团队让模型预测干净的数据，通过调整损失函数实现。

最后一步，作者又让模型直接预测原始图像，而不是经过PCA编码的潜空间。这涉及到在PCA空间中计算残差，并在损失函数中对PCA重建误差进行加权。

最终，整个模型的工作流程就非常简单了：

输入是一张有噪声的图片，噪声添加在PCA潜空间里。输出是原始的干净图片。

经过这一系列改造，团队发现：

• 模型的表示能力主要来源于去噪过程，而非扩散过程。
• 低维潜在空间非常关键，允许模型学习到数据的压缩表示。

并且提出只有很少的现代组件对学习良好的表示至关重要，而许多其他组件是非必要的。

经过改造简化后的模型称为l-DAE，在自监督学习任务中表现出竞争力，同时在结构上更接近于传统的DAE（去噪自编码器）。

在实验中，MAE和l-DAE两种自编码器方法表现都超过了MoCo v3，在ViT-B（86m参数）上，l-DAE表现与MAE相当，但在更大规模的ViT-L（304M）上还所有落后。

最后，团队在简短的结论中提出：

希望我们的发现能够重新激发对基于去噪方法在当今自监督学习研究领域的兴趣。

两位作者已离开Meta

论文一经公布，几位作者就激情当起了自个儿的首批自来水（doge）。

谢赛宁在推特上转发了别人对I-DAE的分享和讨论，并表示：

越来越多研究表明，扩散模型也可以作为效果优秀的特征提取器来使用。

作者刘壮也在一旁开麦，表示扩散模型这东西，不仅仅能做生成相关的工作：

还有开发者惊讶地发现，居然还有人使用主成分分析，可有一阵子不见了。

不过，大神们的工作刚放出来嘛，还有待大家多多传播和具体感知。

除这项最新研究涉及了MAE外，何恺明在提出代表作MAE后，还在此基础上还发表了一系列研究。

例如，提出了一个用来训练CLIP的快速、简单且有效的方法FLIP（Fast Language-Image Pre-training），对CLIP架构引入类似于MAE的思路。

只在模型的极简结构上，施加了一个简单的mask，就让新模型的速度快了3.7倍，同时，性能还可以做到不降反升。

此外，还提出了不对ViT引入分层设计，只用普通ViT就能搞定目标检测的ViTDet。

他们使用普通ViT作为骨干网络，基于MAE方法进行预训练，由此得到的ViTDet能与之前所有基于分层骨干网络的先进方法竞争。

其他学者同样将MAE方法用在了众多项目之中，MAE开山论文目前谷歌学术引用量达4174次。

而何恺明近期的一次公开露面，是在香港中文大学参加了一个讲座，讲道“自己做科研也会emo”。

同时，他在讲座上回答了校友们提出的关于大模型、AI for Science等相关的诸多问题。

其中干货满满，在现场的多半排了很久的长队才一睹大神真容，不在现场的则像追剧似的搜罗网友发在网上的演讲片段视频。

一作陈鑫磊， 目前是Meta FAIR实验室的研究科学家，也是浙大校友，研究兴趣集中于预训练，特别是有自监督或是多模态视觉表示的预训练。

作者刘壮，同样是Meta FAIR实验室的研究科学家，同时也是清华姚班校友，研究重点为神经网络。

至于另一位合作者谢赛宁这边，他最近从Meta AI离职加入了纽约大学，不过依然还是LeCun的同事。

论文地址：https://arxiv.org/pdf/2401.14404.pdf

参考链接：
[1]https://twitter.com/sainingxie/status/1750741794080407893