最近OpenAI提出的Sora视频生成模型引起了广泛的关注,其背后的扩散模型再一次向我们展示了强大的生成潜力。学术界正积极的探索扩散模型的能力边界,本文介绍一篇来自新加坡国立大学的尤洋团队、UC伯克利和Meta AI Research合作完成的文章。本文另辟蹊径,直接用扩散模型生成神经网络的参数,本文的方法非常简单,仅设计了一个自动编码器和一个标准的隐扩散模型。编码器通过对已训练的网络参数子集提取隐层表示,然后将其送入到扩散模型中进行加噪去噪处理,最后再通过解码器生成新的网络参数子集。作者在多种架构和数据集上进行的实验表明,本文的扩散模型可以以较小的成本生成高质量的神经网络参数。

论文题目: Neural Network Diffusion 论文链接: https://arxiv.org/abs/2402.13144 代码仓库: https://github.com/NUS-HPC-AI-Lab/Neural-Network-Diffusion

值得注意的是,Yann LeCun第一时间在社交媒体上对本文工作进行了转发和评论,他认为使用扩散模型来生成神经网络参数是一个非常有趣的想法。

目前,该工作受到了社区的广泛关注,在推特上的浏览量已经超过446K次,like数也超过了1.2K次, 并且引起了学术界的激烈讨论。

一、引言

扩散过程的思想最早起源于非平衡热力学,在2015年,Sohl-Dickstein等人[1]开始将扩散过程引入到图像生成中,其通过逐步消除输入中的噪声来复原出清晰的图像。 近几年,扩散模型迎来了显著的发展,例如Guided Diffusion向我们证明,扩散模型在图像生成领域拥有超越GAN的能力, 再后来出现的GLIDE、Imagen、DALL E-2和Stable Diffusion均能够生成照片级的逼真图像。OpenAI最新发布的Sora模型更是产生了涌现级的时空序列生成能力

尽管扩散模型在视觉生成中取得了巨大成功,但其在其他领域的潜力尚未得到充分探索。本文作者认为,如果我们仔细观察神经网络的训练过程和扩散模型的扩散过程,可以发现,扩散图像生成与随机梯度下降(SGD)过程存在以下共同点(如上图所示):

(1)神经网络训练和扩散模型的逆过程都可以看作是从随机噪声/初始化到特定分布的过渡

(2)高质量图像和高性能参数均可以通过多次添加噪声退化为简单分布, 例如高斯分布。

传统参数生成方法侧重于从先验和概率建模方面进行探索,例如使用随机神经网络和贝叶斯神经网络,本文引入了一种全新的参数生成方法,称为神经网络扩散(p-diff),其中p代表参数(parameters)。虽然参数生成与传统的视觉生成任务有本质上的区别,但本文提出的p-diff仍然能够良好的实现特定参数分布的生成

二、本文方法

本文的p-diff框架构成非常简单,其目的是从无序的随机噪声中生成高性能的模型参数,如下图所示,p-diff的pipeline由两个过程构成,分别是参数自动编码器和参数生成。 首先给定一组已训练的模型,从中选择一部分参数子集并进行展平得到一维向量。随后引入一个自编码器从这些向量中提取隐层表示,并设置一个解码器从隐层特征出发重建原始参数。在参数生成过程中,作者引入了一个标准隐扩散模型(latent diffusion model,LDM)来从随机噪声生成隐层表示

2.1 参数自动编码器

p-diff框架的第一步是对参数自编码器进行训练,本文作者收集了一系列经过SGD优化器更新得到的模型参数 ,其中 是训练样本的数量对于一些已经在大规模数据集上训练过的大型架构,考虑到从头开始训练它们的成本,作者仅对预训练模型的参数子集进行微调,并将微调后的参数保存为训练样本

然后,作者将这些参数 展平为一维向量 ,其中 , 是子集参数数量。随后通过自编码器来重建这些参数 ,为了增强自编码器的鲁棒性和泛化性,作者同时在输入参数和隐层表示中加入随机噪声增强, 整体的编码和解码过程可以表示为:

其中 和 分别表示编码器和解码器。 表示隐层表示, 和 表示添加到输入参数 和隐层表示 中的随机噪声, 是重构得到的参数。与正常的自编码器训练目标相同,p-diff也通过最小化之间的均方误差(MSE)损失进行优化

2.2 参数生成

在进行参数生成时,最直接的做法是使用扩散模型LDM进行处理,但是这个操作的内存成本太高,尤其是当遇到较大规模的模型时,基于这样的考虑,作者选择将降维后的隐层特征作为操作对象, 对于从参数自编码器中提取得到的 ,我们使用DDPM[2]中的策略进行优化:

其中 为添加的高斯噪声, 为去噪网络。在完成对参数生成的训练后,就可以直接将随机噪声输入到逆向过程和经过训练的解码器中来产生一组新的高性能参数,新参数可以与剩余的模型参数合并起来得到新的模型。 由于神经网络参数与图像像素数据不同,其大多没有空间相关性,因此作者在参数自动编码器和参数生成过程中使用一维卷积替换了二维卷积。

三、实验效果

3.1 对比实验和消融研究

本文的实验在多种标准的数据集和网络架构上进行,其中数据集包括MNIST、CIFAR-10/100、ImageNet-1K、 STL-10、Flowers、Pets 和 F-101,网络结构包括 ResNet-18/50、ViT-Tiny/Base 和 ConvNeXt-T/B,以保证实验的普适性和完善性。在实际训练时,自编码器和隐扩散模型均使用4 层 1D CNN 的编码器和解码器。在性能对比时,作者将原始模型中最好的验证精度标识为“original”,将原始模型的平均权重集成标识为“ensemble”。

上表展示了本文方法在8个数据集6个网络架构上的性能对比结果,每组数据的表示格式为“original / ensemble / p-diff”,从表中可以得到两个观察结果:

(1)在大多数情况下,p-diff取得了与两个基线相似或更好的结果。这表明p-diff可以有效地学习高性能参数的分布,并从随机噪声中生成优秀的模型

(2)p-diff在各种数据集上始终表现良好,这表明本文方法具有良好的普适性

上表展示了作者进行的消融实验结果,表(a)展示了不同训练数据规模(原始模型的数量)对p-diff性能的影响,可以发现,训练数据多少对模型最佳结果的性能差距很小

为了评估p-diff在不同网络层数位置的生成效果,作者将p-diff应用到三种不同深度的BN层中,如上表(b)中所示,作者发现,p-diff在所有深度的BN层中均获得了比原始模型更好的性能。 此外,作者还对训练自编码器时引入的噪声增强策略进行了消融研究,消融结果如上图(c)所示。

为了探究p-diff的背后的内在模式,作者设置了三个随机种子从头训练 ResNet-18 模型,并对其中不同层(Conv.-layer2和FC-layer18)的参数进行可视化,如上图所示,这些层之间确实存在特定的参数模式,通过对这些模式进行学习,p-diff得以生成高性能的神经网络参数。

3.2 p-diff是仅仅记忆了参数吗

虽然p-diff在模型精度方面已具有一定的优势,但是这些结果其是否只是通过记住训练集中原始模型的参数实现的。 本文作者希望p-diff能够真正生成一些与原始模型分布不同的新参数。为了验证这一点,作者首先提出了一个量化指标来衡量两个模型参数之间的相似性, 具体通过计算两个模型在错误预测结果上的IoU来实现。IoU 越高表明两个模型的预测之间的相似性越大。

上图(a)展示了四种相似性比较:1)原始模型之间的相似性; 2)p-diff模型之间的相似性; 3)原始模型和p-diff模型之间的相似性; 4) 原始模型和 p-diff 模型之间的最大相似度。可以看到,p-diff生成的模型之间的差异远大于原始模型之间的差异。这表明p-diff可以生成与其训练数据(即原始模型)表现不同的新参数。

此外,作者还将p-diff与微调和添加噪声后的模型进行了对比,上图(b)的结果表明,仅进行微调和添加噪声的模型性能很难超越原始模型。 此外,微调或添加噪声的模型与原始模型之间的相似度非常高,这表明这两种操作无法获得新颖高性能的模型

图(c)展示了原始模型与p-diff生成模型的隐层表示的分布情况(使用t-SNE算法进行降维),可以看到p-diff可以生成与原始模型分布相差较大的新参数,而普通的添加噪声的方法只能在原始模型的隐层表示周围进行插值

四、总结

本文提出了一种使用扩散模型生成神经网络参数的框架p-diff,p-diff可以从传统的SGD优化器得到的参数出发,生成具有较大差异的高性能网络参数, 这一点与先前训练模型的过程有很大的差别,因而其能够始终相比原始训练数据实现相似甚至更强的性能。 此外,本文作者也提到,由于网络参数是与图像视频性质不同的信号,因此必须谨慎考虑处理这种区别,虽然p-diff现在仍然无法生成现代规模的大模型参数,但是在训练算力焦虑笼罩的今天,这种想法非常振奋人心。我们期望更加优秀的扩散模型参数生成框架的出现。

参考

[1] Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML. PMLR, 2015.

[2] Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. NeurIPS, 33, 2020.

Illustration From IconScout By Delesign Graphics

-The End-

扫码观看!

本周上新!

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区(www.techbeat.net)。社区上线500+期talk视频,3000+篇技术干货文章,方向覆盖CV/NLP/ML/Robotis等;每月定期举办顶会及其他线上交流活动,不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台,希望为AI人才打造更专业的服务和体验,加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章,并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向,对用户启发更大的文章,做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信(chemn493)投稿,沟通投稿详情;还可以关注“将门创投”公众号,后台回复“投稿”二字,获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域新型创投机构,也是北京市标杆型孵化器。 公司致力于通过连接技术与商业,发掘和培育具有全球影响力的科技创新企业,推动企业创新发展与产业升级。

将门成立于2015年底,创始团队由微软创投在中国的创始团队原班人马构建而成,曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业,不仅想获得投资,还希望获得一系列持续性、有价值的投后服务,欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角,把文章分享到朋友圈