可逆ResNet:极致的暴力美学
作者丨苏剑林
单位丨广州火焰信息科技有限公司
研究方向丨NLP,神经网络
个人主页丨kexue.fm
今天我们来介绍一个非常“暴力”的模型:可逆 ResNet。
为什么一个模型可以可以用“暴力”来形容呢?当然是因为它确实非常暴力:它综合了很多数学技巧,活生生地(在一定约束下)把常规的 ResNet 模型搞成了可逆的!
▲ 标准ResNet与可逆ResNet对比图。可逆ResNet允许信息无损可逆流动,而标准ResNet在某处则存在“坍缩”现象。
模型出自 Invertible Residual Networks,之前在机器之心也报导过。在这篇文章中,我们来简单欣赏一下它的原理和内容。
可逆模型的点滴
为什么要研究可逆 ResNet 模型?它有什么好处?以前没有人研究过吗?
可逆的好处
可逆意味着什么?
意味着它是信息无损的,意味着它或许可以用来做更好的分类网络,意味着可以直接用最大似然来做生成模型,而且得益于 ResNet 强大的能力,意味着它可能有着比之前的 Glow 模型更好的表现。
总而言之,如果一个模型是可逆的,可逆的成本不高而且拟合能力强,那么它就有很广的用途(分类、密度估计和生成任务,等等)。
而本文要介绍的可逆 ResNet 基本上满足这几点要求,它可逆起来比较简单,而且基本上不改变 ResNet 的拟合能力。因此,我认为它称得上是“美”的模型。
旧模型的局限
可逆模型的研究由来已久,它们被称为“流模型(flow-based model)”,代表模型有 NICE、RealNVP 和 Glow,笔者曾撰文介绍过它们,另外还有一些自回归流模型。除了用来做生成模型,用可逆模型来做分类任务的也有研究,代表作是 RevNet [1] 和 i-RevNet [2]。
这些流模型的设计思路基本上都是一样的:通过比较巧妙的设计,使得模型每一层的逆变换比较简单,而且雅可比矩阵是一个三角阵,从而雅可比行列式很容易计算。
这样的模型在理论上很优雅漂亮,但是有一个很严重的问题:由于必须保证逆变换简单和雅可比行列式容易计算,那么每一层的非线性变换能力都很弱。事实上像 Glow 这样的模型,每一层只有一半的变量被变换,所以为了保证充分的拟合能力,模型就必须堆得非常深(比如 256 的人脸生成,Glow 模型堆了大概 600 个卷积层,两亿参数量),计算量非常大。
硬“杠”残差模型
而这一次的可逆 ResNet 跟以往的流模型不一样,它就是在普通的 ResNet 结构基础上加了一些约束,就使得模型成为可逆的,实际上依然保留了 ResNet 的基本结构和大部分的拟合能力。这样一来,以往我们在 ResNet 的设计经验基本上还可以用,而且模型不再需要像 Glow 那样拼命堆卷积了。
当然,这样做是有代价的,因为没有特别的设计,所以我们需要比较暴力的方法才能获得它的逆函数和雅可比行列式。所以,这次的可逆 ResNet,很美,但也很暴力,称得上是“极致的暴力美学”。
可逆“三要素”
ResNet 模型的基本模块就是:
也就是说本来想用一个神经网络拟合 y 的,现在变成了用神经网络拟合 y−x 了,其中 x,y 都是向量(张量)。这样做的好处是梯度不容易消失,能训练深层网络,让信息多通道传输,等等。可逆的意思就是 x+g(x) 是一个一一映射,也就是说每个 x 只有一个 y 与之对应,反过来也是,换言之我们理论中可以从中解出反函数 x=h(y) 来。
背景就不多说了,但是要说明一点,我们在分类问题中用的广义上的 ResNet,是允许通过 1×1 卷积改变维度的,但这里的 ResNet 指的是不改变维度的 ResNet,也就是说 x,y 的大小保持一样。
对于一个号称“可逆”的模型,必须要回答三个问题:
什么时候可逆?
逆函数是什么?
雅可比行列式怎么算?
从难度上来看,这三个问题是层层递进的。当然,如果你只关心做分类问题,那么事实上保证第一点就行了;如果你关心重构图片,那么就需要第二点;如果你还想像 Glow 那样用最大似然来训练生成模型,那么就需要第三点。
下面按照原论文的思路,逐一解决这三个问题(三道“硬菜”),来一场暴力盛宴。
什么时候可逆?
第一道硬菜是三道硬菜中相对容易啃的,当然只是“相对”,事实上也都用到了泛函分析的一些基础知识了。
因为 (1) 是 ResNet 的基本模块,所以我们只需要保证每个模块都可逆就行了。而 (1) 可逆的一个充分条件是:
其中:
是函数 g 的 Lipschitz 范数。也就是说,g 的 Lipschitz 范数小于 1,就能保证 (1) 可逆了。
那什么时候 g 的 Lipschitz 范数才会小于 1 呢?因为 g 是神经网络,卷积和全连接都无妨,神经网络是由矩阵运算和激活函数组合而成的,即代表结构是:
那么由链式法则,“g 的 Lipschitz 范数小于 1”的一个充分条件是“Activation 的 Lipschitz 范数不超过 1”且“Wx+b 的 Lipschitz 范数小于 1”。而 Activation 只是个标量函数,“Lipschitz 范数不超过 1”意味着导数不超过 1 就行了,目前我们常用的激活函数(sigmoid、tanh、relu、elu、swish 等)都满足,所以这一点不用管它;而“Wx+b 的 Lipschitz 范数小于 1”,意味着矩阵 W 的 Lipschitz 范数小于 1。
矩阵 W 的 Lipschitz 范数其实也就是“谱范数”,记为 Lip(W) 或
对模型 g 的所有核权重 W 做谱归一化,然后乘上一个大于 0 小于 1 的系数 c(即
逆函数是什么?
为什么这样就可逆了?证明过程可以直接回答第二个问题,也就是说,我们直接把逆函数求出来,然后就知道什么条件下可逆了。
假如 y=x+g(x) 是可逆的,那么我们要想办法求出逆函数 x=h(y),这其实就是解一个非线性方程组。简单起见,我们考虑迭代:
显然,迭代序列 {Xn} 是跟 y 有关的,而一旦 {Xn} 收敛到某个固定的函数:
那么我们就有
换句话说,如果迭代 (5) 收敛,那么收敛的结果就是 x+g(x) 的逆函数。所以我们只需要搞清楚 (5) 什么时候收敛。这时候前面的条件 Lip(g)<1 就能用上了,我们有:
所以:
可以看到,
附:单纯指出
对于任意正整数 k,我们继续考虑
可以看到我们得到了
至此,我们终于证明了 {Xn} 的收敛性。顺便一提的是,在 (9) 中,取 k→∞,我们得到:
也就是说,这个迭代算法的收敛速度跟正比于
说宏大一点,其实这就是泛函分析中的“巴拿赫不动点定理”,又称“压缩映射定理”(因为 Lip(g) 小于 1,所以 g 被称为一个压缩映射)。
这样一来,我们已经回答了为什么 Lip(g)<1 就有 x+g(x) 可逆了,同时已经给出了求逆函数的方法——就是通过 (5) 迭代到足够的精度:
当做好归一化操作使得 x+g(x) 可逆后,它的逆函数为
终于,我们啃下了第二道硬菜。
雅可比行列式怎么算?
下面来到三个问题中最“硬核”的一个问题:雅可比行列式怎么算?为了解决它,作者综合了数学分析、矩阵论、概率论、统计采样等多方面数学工具,堪称“暴力之最”、“硬菜之最”。
首先,为什么要算雅可比行列式?前面已经说了,只有做生成模型时才有这个必要,具体细节请参考笔者最早的对流模型的介绍《细水长flow之NICE:流模型的基本概念与实现》。接着,我们知道雅可比行列式就是雅可比矩阵的行列式,所以要把雅可比矩阵算出来:
再次提醒,虽然我偷懒没有加粗,但这里的 g 输出是一个向量,x 也是一个向量,∂g/∂x 实际上就是输入和输出两两配对求偏导数,结果是一个矩阵(雅可比矩阵)。
然后,雅可比行列式就是
最后一个恒等号,是因为 det(I+Jg) 一定是正数,所以可以去掉绝对值。这是可以证明的,但只是细节部分,我们就不纠结了,读者自行去看作者提供的参考文献吧。
然后呢?直接按定义来计算雅可比行列式?不行,因为这样子计算量实在是太大了,而且反向传播的时候,还要算行列式的导数,那就更复杂了。作者们想出了繁琐但有效的解决方法,利用恒等式(参考《恒等式 det(exp(A)) = exp(Tr(A)) 赏析》[3]):
我们得到:
如果能求出 ln(I+Jg) 来,然后求迹(trace,对角线元素之和)就行了。怎么求 ln(I+Jg) 呢?还是参考《恒等式 det(exp(A)) = exp(Tr(A)) 赏析》[3],暴力展开:
注意这个级数收敛的条件是
现在 ln(I+Jg) 变成了一个无穷级数,如果截断 n 项,那么误差也正比于
问题解决了吗?上式需要我们去计算
假设 p(u) 是一个多元概率分布,其均值为 0、协方差为单位矩阵(显然标准正态分布符合要求),那么对于任意矩阵 A,我们有:
利用“均值为 0、协方差为单位矩阵”这个性质,直接按定义就可以证明上式了,并不困难。然后,作者提出了一个显得“既无赖又合理”的方法:对于每次迭代,我只从 p(u) 中随机选两个向量 u1,u2 出来,然后认为
读者可能就有意见了,不是要对所有向量求平均吗?只随机挑两个就行了?其实还真的是可以了,理由如下:
1. 我们优化都是基于随机梯度下降的,本来就带有误差,而只随机挑两个也有误差,而每步迭代都重新挑不同的 u1,u2,在一定程度上就能抵消误差;
2. 更重要的原因是,我们要算雅可比行列式的对数,只是用它来做一个额外的 loss,来保证模型不会坍缩,简单来讲,可以将它看称一个正则项,而既然是一个正则项,有点误差也无妨。
所以,
所以不需要把
所以,最终可以总结为:将雅可比矩阵做对数级数展开 (15),然后将行列式计算转化为迹的计算 (16),并且利用概率采样的方式 (18),就可以最高效地计算雅可比行列式。
纵观实验效果
其实笔者一开始是就被“可逆 ResNet”这么美好的构思吸引过来的,但是看到这里,我发现我也怂了,这绝对对得起“硬杠 ResNet”的评价呀。本来想对照着好歹实现个 mnist 的生成来玩玩,后来确认有这么多技巧,如此之暴力,我也放弃了。
所以,还是来看看原论文的实验结果好了。
Toy数据集
首先是一个人造的 Toy 数据集,也就是构造一些有规律的随机点来,然后用生成模型去拟合它的分布,GAN 中也常做这样的实验。
从下图可以看到可逆 ResNet 的效果比 Glow 好得多,归根结底,我觉得是因为可逆 ResNet 是一个很对称的、没有什么偏置的模型,而 Glow 则是有偏的,它需要我们以某种方式打乱输入,然后对半切分,对半之后两部分的运算是不一样的,这就带来了不对称。
▲ 可逆ResNet实验:Toy数据集
类任务实验
一开始我们就说了,既然是 ResNet 的一种,最基本的用途就是分类了。下表表明用可逆 ResNet 做分类任务,效果也是很不错的,Lipschitz 约束的存在不会明显影响分类结果(表中的 c 就是 Lip(g))。
▲ 可逆ResNet实验:分类效果
生成模型实验
其实流模型系列在生成复杂图片方面还远比不上 GAN,但相互之间的定量对比倒没有问题。下图也表明可逆 ResNet 作为一个基于流的生成模型方面也是很优秀的。
▲ 可逆ResNet实验:生成模型效果
终于可以收工了
这样活生生地取硬杠 ResNet 可真是一件苦力活,我就单纯的解读就已经这么累了,真佩服作者的数学功底。当然,最后作者终究还是成功了,想必成功的喜悦也是很丰盛的。总的来说,整个工作很暴力,但细细品味之下并没有什么违和感,反倒是有一种浑然的美感在里边,并非简单的堆砌数学公式。
唯一的问题是,整个“硬杠”的流程还是挺复杂的,因此要推广使用还是得有好的封装,而这往往就让很多人望而却步了。还有一个问题,就是流模型系列为了保证可逆,自然是不能降维的,但不降维必然就导致计算量大,这是个矛盾的地方。
一个有趣的想法是,对于降维的情形,能不能搞个类似矩阵的“伪逆”那样的模型,来达到类似可逆ResNet的效果呢?非方阵也可以搞个行列式出来的(比如《再谈非方阵的行列式》[4]),因此降维变换也应该能够搞个雅可比行列式的对数出来。貌似不少地方是可以推广过来的。
流模型后面的方向究竟如何呢?让我们拭目以待。
参考文献
[1] Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse. The Reversible Residual Network: Backpropagation Without Storing Activations. arXiv preprint arXiv:1707.04585, 2017. 7
[2] Jörn-Henrik Jacobsen, Arnold Smeulders, Edouard Oyallon. i-RevNet: Deep Invertible Networks. In Proceedings of the International Conference on Learning Representations (ICLR), 2018.
[3] https://kexue.fm/archives/6377
[4] https://kexue.fm/archives/6096
点击以下标题查看作者其他文章:
让你的论文被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢? 答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学习心得或技术干货。我们的目的只有一个,让知识真正流动起来。
📝 来稿标准:
• 稿件确系个人原创作品,来稿需注明作者个人信息(姓名+学校/工作单位+学历/职位+研究方向)
• 如果文章并非首发,请在投稿时提醒并附上所有已发布链接
• PaperWeekly 默认每篇文章都是首发,均会添加“原创”标志
📬 投稿邮箱:
• 投稿邮箱:hr@paperweekly.site
• 所有文章配图,请单独在附件中发送
• 请留下即时联系方式(微信或手机),以便我们在编辑发布时和作者沟通
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
关于PaperWeekly
PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域,欢迎在公众号后台点击「交流群」,小助手将把你带入 PaperWeekly 的交流群里。
▽ 点击 | 阅读原文 | 查看作者博客