0%

最近,我在网上看到了一个专门介绍深度学习源代码的网站(https://nn.labml.ai/index.html ),各类流行网络架构一应俱全。这个网站最吸引我的,是它的双列式结构。如下所示,每个页面右边显示代码,左边显示每段代码对应的注释。

一直以来,我在博客中介绍代码时,都是先写描述文字,再贴一段代码。这种方式对作者和读者来说都十分低效。受到上面那个网站的启发,我决定以后也采用这种方式介绍源代码。很可惜,上面那个网站是由 labml.ai 这个组织维护的,似乎并没有提供开源的、可定制的注释网站搭建方式。为此,我找到了一个替代品:Pycco。它可以方便地为 Python 代码生成双列带注释的静态页面。

Pycco 安装与使用

Pycco 是页面生成工具 Docco 的 Python 实现版。它们都是「快而脏」的文档生成器:源代码只有几百行,没有多余封装,直观而暴力地生成 HTML 页面的所有内容。它们为初学者只提供了一个文档页面,这个文档页面讲的是它们的源代码,且文档页面就是由它们本身创建出来的。Pycco 的官方文档页面如下所示。一开始介绍完 Pycco 的背景信息和安装方式后,文档就会直接开始介绍源代码。

根据文档的指示,我们可以通过下面的命令在 pip 里安装 Pycco:

1
pip install pycco

在使用 Pycco 时,我们完全不用去理解其源代码的意思,只需要准备一个带注释的 Python 源文件就行。Pycco 提供了一键式命令,帮我们把源文件里的注释和代码分离,并生成左侧为注释,右侧为代码的静态文档页面。这里的注释既包括了 # 开头的单行注释,也包括了 """ """" 包裹的多行注释。注意,三个单引号构成的多行注释 ''' ''' 不会被该工具识别。

我们来看一个示例。假设我们有下面一个名为 hello.py 的 Python 源文件。

1
2
3
4
5
6
7
8
9
10
11
12
13
"""
Hello
World
"""
import torch
import pycco

# Hello
print("hello world")

"""
End of file
"""

我们用下面的命令为其生成文档页面。

1
pycco hello.py

页面会生成在目录下的 docs 子目录里,子目录包含一个 hello.html 文件和一个 pycco.css 文件。我们可以用浏览器打开 .html 网页,看到下图所示内容。

注释默认使用 Markdown 语法。如果平时就习惯用 Markdown 写博客的话可以无缝切换。但是,默认情况下网页无法渲染公式。

页面跳转

除了普通的 Markdown 语法外,Pycco 还支持一个特殊的功能:文档跳转。我们可以把文件名写在 [[ ]] 内,实现源文件内部或源文件间的跳转。特别地,我们还可以在页面某处打上标签,并跳转到某页面的标签处。以下是一个示例。为了给「注释」加上注释,我别出心裁地把注释写进了代码部分。

让页面渲染 LaTeX 公式

刚才我们发现,目前的 Pycco 页面并不支持公式渲染。而在解释深度学习代码时,很多时候不得不用到公式。因此,我决定给 Pycco 加上渲染公式的功能。

Pycco 这种直观暴力的实现方法让网页开发者能够快速地修改页面生成逻辑。然而,我已经把 HTML 的知识快忘光了,配不上「网页开发者」这个名号。因此,我让 ChatGPT o1 来帮我开发这一功能。

经指导,我认识了 MathJax 这个在网页上渲染公式的工具。只需要在 HTML 的 head 里导入一些包,网页就可以自动识别单行公式 和多行公式 $$$ $$$。我不记得 head 是什么了,但大概能猜到这个是一个相当于声明全局变量的语句块。

我在 pycco_resources\__init__.py 文件里找到了设置 head 的地方。这个文件提供了生成网页的模板,包括了写死的 CSS 文件和部分 HTML 代码。打开这个文件的最快方式是在某处写 import pycco_resources,之后用 IDE 的代码跳转找到这个包在本地的位置。

我们在该文件下图所示的位置插入和 MathJax 相关的代码。

1
2
3
4
5
6
7
8
<script>
MathJax = {
tex: {
inlineMath: [ ["$","$"] ]
}
};
</script>
<script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-chtml.js"></script>

修改模板后,重新生成网页,可以发现现在公式能够正确渲染了。

高亮新增/删除代码

在以后讲解代码的时候,我想高亮新增的或者删去的代码,就像 GitHub 里显示代码的更改一样。由于 Pycco 非常容易做二次开发,我又去请教 ChatGPT o1 该如何实现这样的功能。

经指导,我了解到代码高亮可以通过设置背景颜色 style 来实现。为此,我需要做两件事:

  1. 新增有关背景颜色的 CSS style
  2. 想办法指定要高亮的代码块

第一件事很简单,只需要打开模板文件 pycco_resources\__init__.py,添加背景颜色 style 即可。我添加了两种背景颜色,分别表示删去的和新增的代码块。我是在 VSCode 里找到想要的颜色值的。随便打开一个 CSS 文件,输入一个大概的颜色值,VSCode 就会弹出一个选择颜色的小窗口。

1
2
3
4
5
6
.highlighted-line-1 {
background-color: #fa53208c;
}
.highlighted-line-2 {
background-color: #68fc5d;
}

第二件事就比较难了。我需要先看懂目前高亮源代码的逻辑,再在其基础之上添加背景高亮的功能。Pycco 的主函数在文件 pycco/main.py 里,我们可以用导入 Python 包的方式快速找到这份源文件。原来的高亮关键字的逻辑如下,我在其中加入了一些代码用于输出中间结果。

函数的主要输入是列表 sections,函数的输出存储在列表项 section"docs_html", "code_html" 字段。

sections 内容如下,可见其中是过了解析器的注释块和代码块。每一段注释都会对应一段代码。

1
[{'docs_text': '=== BeGin ===\r\n\n', 'code_text': 'import torch\r\nimport pycco\r\n\r'}, {'docs_text': '# Hello\r\n## World\r\nThis is a **Python** [code](https://www.python.org/).\r\nTry formula $\\epsilon$ .\r\n$$\r\n3 = 1 + 2\r\n$$\r\n\n\n', 'code_text': 'print("hello world")\r\n\r'}, {'docs_text': 'End of file\r\n\n', 'code_text': '\n'}]

函数的输出是完整 HTML 文件的一部分。

1
2
3
4
5
6
7
8
9
10
11
12
13
<h3><span id="begin" href="begin"> BeGin </span></h3>
<div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">pycco</span></pre></div>
<h1>Hello</h1>
<h2>World</h2>
<p>This is a <strong>Python</strong> <a href="https://www.python.org/">code</a>.
Try formula $\epsilon$ .
$$
3 = 1 + 2
$$</p>
<div class="highlight"><pre><span class="nb">print</span><span class="p">(</span><span class="s2">&quot;hello world&quot;</span><span class="p">)</span></pre></div>
<p>End of file</p>
<div class="highlight"><pre></pre></div>

那么,我的目标就是修改这个文件。我需要先根据输入的原始代码块,判断这段代码是否要高亮,再修改 HTML 代码块的内容。

首先,从用户的角度考虑,应该怎么指定要背景高亮的代码呢?既然现在代码被拆成了一块一块的,且每块代码对应一段注释,我决定用一些特殊注释来高亮一整块代码,就像前面的设置标签和跳转标签的特殊注释一样。我加入了以下判断:如果注释前几个字符是 '===ADD===''===DEL===',就用对应的颜色高亮这段代码。

判断了是否要高亮后,我还需要做对应的修改。我不仅要在 HTML 代码块里高亮代码,还需要把注释块里的特殊命令删掉。通过观察相关代码,我忽然回忆起了 HTML 的部分实现原理:背景高亮就是把一段 HTML 字符串封进一个带有背景高亮样式的标签 <div></div> 里。那剩下的删除注释也很简单,只需要对字符串做一点简单操作就行了。代码修改过程及使用示例如下所示。

1
2
3
4
5
6
7
8
9
10
11
12
# Add following code

delete_background_start = '<div class="highlighted-line-1">'
add_background_start = '<div class="highlighted-line-2">'
background_end = '</div>'

if section["docs_text"].startswith('===ADD==='):
section["docs_text"] = section["docs_text"][9:]
section["code_html"] = add_background_start + section["code_html"] + background_end
elif section["docs_text"].startswith('===DEL==='):
section["docs_text"] = section["docs_text"][9:]
section["code_html"] = delete_background_start + section["code_html"] + background_end

修正多行注释缩进

在写这篇博文的时候,我又发现 Pycco 的一个 bug:在多行注释带缩进时,HTML 页面上的注释也会缩进,导致页面布局完全乱套。

我找到了 pycco/main.py 里处理多行注释的代码。我把旧代码全部删掉,换了一种暴力的字符串处理方式:删除多行注释左侧所有空格及制表符。

1
docs_text += line.lstrip('\t ') + '\n'

改完之后 bug 就修好了。

总结

在这篇文章中,我介绍了 Pycco 这个为 Python 代码生成双列注释-代码静态页面的 Python 工具,并分享了我对其的三个修改:支持 MathJax 公式、支持背景高亮、修复多行注释缩进 bug。相比介绍开源项目的 readthedocs 等文档构建工具,Pycco 能够为每行 Python 代码写注释,尤其适合详细讲解深度学习代码。对于熟练的网页开发者,Pycco 的源代码很短,改起来很方便;但另一方面,Pycco 似乎比较小众,多年都没人维护,可能有各种 bug。主要推荐想在个人博客里讲解代码的朋友使用这个工具。我以后都会用它来讲解代码。

ChatGPT 小插曲

ChatGPT 在指导我添加 MathJax 时,先叫我下载 python-markdown-math 这个 pip 包,让 Python 脚本里的 Markdown 解析器能够解析公式。但我在实践中发现不解析 $$$$ 符号,在正确导入 MathJax 包后,网页依然能够正确渲染。

ChatGPT 一开始提供的 MathJax 导入代码是 MathJax V3 的。经过 python-markdown-math 解析后,双行注释被转换成了 <script type="math/tex; mode=display"></script> 包裹的内容。这些公式没法成功渲染。我把添加双行公式的详细过程提供给 ChatGPT 后,它分析出这种 <script> 的写法是 MathJax V2 的,让我把 head 里的包改成 MathJax V2。修改过后,双行公式果然渲染成功了。然而,如上一段所述,我发现哪怕不用 python-markdown-math 也行。于是我又用回了 MathJax V3。可见 ChatGPT 有功有过,给我制造了一些麻烦,也算是能够帮我解决。

在求助 ChatGPT 添加背景高亮功能时,它非常出色地分析了问题:先理解 Pycco 生成 HTML 的过程,再决定用户如何指定要高亮的内容,最后根据不同的方案做实现。但它给的三个信息输入方案都不尽如人意:(1) 在命令行里指定要高亮的行;(2) 在每行要高亮的代码前都加特殊注释;(3) 用另一个配置文件来指定要高亮的行。这些方法的用户体验都不好,且难于实现。当然,由于我没有把所有源代码给它,它或许没有发挥完整的实力。最后我还是靠自己理解了 Pycco 的页面生成逻辑,选了一种好实现且好用的方案。

我从来没有用 ChatGPT 做过特别复杂的编程任务,这次稍微体会了一下。感觉它的能力有时候会出乎意料地强(比如发现 HTML 的风格是 MathJax V2 的),但不见得次次都靠谱。它的能力主要还是基于大量数据的「搜索」,加上一点浅显的思考。只是因为看过的数据过多,所以很多时候能从历史经验中直接给出很好的答案。当然,作为语言模型,它对于语言的理解能力极强,哪怕描述不完全也能准确理解用户的意图。做为使用者,我们要在用之前就想好能从 ChatGPT 那里获取到哪些帮助,并提供充足的信息,而不能全盘接受它的输出,就和我们日常和别人交流一样。

年底难得有闲暇,我打算首次写一下年终回顾。写这样的文章总是很害羞,我现在不好意思去看我以前写的「随笔」类文章了。但是,今天,我就是有一股冲劲,想要实现心中那小小的想法,不想被生活的洪荒冲走我那鲜活的热情。不去思考为谁而写,不去思考是否有人会看,不去纠结于遣词造句。只求能够纯粹地表达。

人的思考有时非常奇怪。有时会对未来产生过分的期待,期待过后又是更深的失落。小的时候,盼望长大。可自从大学毕业以后,见过了周围人的种种言谈,我又害怕起了变老。明明二十多岁,正值人生中精力最鼎盛的时期,一向乐观的我也时常害怕起了年龄变大,害怕像周围人一样慢慢失去对未来的期盼。就是在这样一种心理下,我在 2024 年却碰到很多令人惊喜的事,让我继续有了向前看的动力。


先说游戏吧。我一直觉得这几年新游戏都没什么创新,但从今年年初开始,我见到了非常多有意思的游戏。「小丑牌」将自走棋的经济系统与爬塔类游戏结合,创造出了一款机制简明却不失趣味与深度的游戏,挖掘出了 Roguelike 类游戏好玩的本质。《背包乱斗》拓展了以前背包类游戏的机制,通过 PVP 自走棋的游戏形式让游戏自带了平衡性与挑战性,从而让背包类游戏的魅力得以绽放。但成也自走棋败也自走棋,做成 PVP 后,游戏的趣味性必然会大打折扣。作者也没能深刻理解自走棋类游戏该如何设计,不去思考怎么做平衡和简化游戏机制,不停地试图加新内容。游戏一下就玩腻了,非常可惜。年中出的《动物井》和《海天岛传奇》我没玩,但我很高兴能见到这些优秀的解谜类游戏。《风暴之城》是今年最令我惊喜的游戏。我从小就很喜欢玩《主题医院》、《龙之崛起》这类模拟经营游戏,但这些游戏的特点是趣味性大于挑战性,战役的每一关的玩法都大同小异,很容易玩腻。而《风暴之城》把近年来的 Roguelike 设计思路引入了模拟经营游戏,并精心设计了会让玩家游戏失败的「危机」,玩游戏时能体会到玩即时战略游戏的那种刺激感。最终,《风暴之城》兼具了 Roguelike、 模拟经营、即时战略等多种游戏的优点,对于喜欢二十多年前的游戏的老玩家来说非常新鲜又亲切。在我心中,《风暴之城》已经代替《英雄无敌3》,成为了综合游戏体验最佳的游戏。另一个我今年才挖掘出来的游戏是《幸运房东》,这个脱胎于爬塔机制的游戏成功开创了「幸运房东类」这一新型游戏。这类游戏操作简单,只需要每次从新的「卡牌」中三选一,最终组建出一套成型构筑。这类游戏把构筑类游戏复杂的「测试构筑」环节(比如暗黑类游戏用当前装备打怪,《杀戮尖塔》里用当前卡组战斗)大幅简化,只单独考验玩家的组建构筑能力,充分把一项游戏机制做到了极致。除了本体外,我还尝试了许多后续「幸运房东类」游戏,最终认识了《轮作法》这个游戏。这个游戏把自走棋玩法融入「幸运房东类」游戏,同时把《炉石传说:酒馆战棋》里的卡牌空位管理机制加入了游戏,让一个有诸多随机机制的游戏充满了思考深度。《轮作法》是我继《岛民》(islander) 以来见过我最像棋类的游戏(信息公开,需要玩家做运算,搜索当前已知最优解),我非常喜欢这类有思考深度的游戏。《流放之路》我也稍微玩了下,但游戏机制没怎么大改,我稍微回味了一下就没玩了。倒是新暗黑类游戏《最后纪元》也让我很惊喜。由于游戏是新出的,攻略很少,我第一次享受了在暗黑类游戏里自己组构筑的快乐。可惜这个游戏重玩体验没有《流放之路》高,第二个赛季我就不太想玩了。去年年底到今年年初我还认识了《战场兄弟》这个游戏,我也没想到这几年来还有类似于《英雄无敌》的以战棋为子玩法的游戏,我很开心能见到这类游戏。

剩下我的日常游戏时间都主要投入在了重玩性强的挑战类游戏上。我从去年下半年开始打《雀魂》,年初被打到道心破碎。后来抽时间稍微玩了一下,总算上了雀豪2,暂时没有继续玩的打算。上半年我又反复提升了《亿万僵尸》的技术,最高难度不暂停的胜率大幅提高。只要我完全认真玩,就很少有过不了的地图。下半年压力很大,我又开始玩《杀戮尖塔》,对游戏的理解越来越深,近期随随便便打出了猎手7连胜。我甚至开始重新捡起了围棋,下的盘数可能快超过了小时候的。但我只和星阵 AI 下,大概能打平 3 段,所以我现在一直在挑战准 4 段。

跟我有类似游戏品味的人太少了,所以十多年来我很少有机会和别人深入交流游戏攻略。而我玩完游戏就忘。所以这两个月我一直在反思我是不是失去了探索新游戏的胆量,只敢玩老游戏了。但看了 Steam 的年终回顾我才发现,原来我今年的游戏之路是这么丰富多彩。我并不想跟随潮流,玩最新最热门的游戏。看到好玩的就玩,不管是几年前出的。顺其自然地玩,才是真正地享受游戏。


今年年中的时候我被直播弹幕引流,去看了「MyGO」。本以为只是消遣,没想到这部动漫彻底改变了我的生活。「MyGO」讲述了几个美少女组乐队的故事。少女们性格各异,因巧合凑到了一起,因矛盾而散开,最终又凭借着人与人之间的牵绊又走回到了一起。动漫对感情的描写非常细腻,且擅长用特色鲜明的场景、台词来展现角色间的冲突。可能是某些场景乍看之下非常滑稽,这部本来仅属于少部分情感细腻的观众的作品被涌入的各式各样的观众玩梗、做二创、引流,形成了别具一格的二创文化。也得益于这些引流,我一个平时不怎么看动漫的人能看到这样一部温馨的作品,我为今年能认识「MyGO」而感到惊喜与感动。

关于「MyGO」的观后感,不管是有关剧内的剧情、剧外的同人创作,还是结合我个人经历得出的评价,我有太多太多想说的话。但是,今天时间有限,我不能快速把它们全部表达出来。我只能想到什么说什么。

作为一部动漫,「MyGO」的表现能力很强,其对于感情的描写甚至能够达到视觉小说的程度(毕竟小说由大量文字组成,更容易表达深刻的思想)。故事中的每个人看似有着各种各样的毛病,但在二次元这个理想的世界里,她们每个人的本质都是纯洁善良的。所以,在矛盾过后,少女们能够把人性中最美好的那些情感展示出来。哪怕大家性格各异,无法简单地用语言达成共识,但是,只要有歌声,有了想要表达自己的冲动,有了想要和他人建立连结的念头,那么演唱过后,那些无法用言语描述的情感会连在一处,并自然地流入每个人的心中。剧中的高松灯并不是我最喜欢的角色,但是她在剧中的做法让我非常感动与向往:作为一名不太敢于他人交流,略有心理障碍的人,她能为了留住和他人的友谊,去忘我地、不计后果地在台上念诗,尽自己一切努力,勇敢地表达着自己。最终,她成功地把队友找回来了。「诗超绊」的名场面,我看一次哭一次。作为极端理性的人,我知道这些事在现实生活中是不可能的。没有耐心的观众,没有愿意追求真挚友情的朋友,自然也没有敢于忘我地表达自己的人。是作者的温柔,让理想的二次元世界里诞生了这样的故事。我知道这样的事是不可能的,不存在的,不应该去梦想的。但是,我羡慕这种能够忘我地追求目标的人。我开始不断反思自己前进的步伐还不够大。我一直说自己有理想,现在我却分不清这到底是真正有想做的事,还是仅仅虚伪地创造一个引诱自己前进的「胡萝卜」了。我不想忘记我的理想。我也想不计后果地去追求!

「MyGO」的二创作品也是盛况空前。歌曲、剪辑、小故事……,催泪向、搞笑向……。看着这么多人都和我有着类似的喜好,都分享着自己的喜好,我感到非常幸福,一度想起了十年前认识「东方」时的感动。如果只有作品本身,「MyGO」对我来说只不过是和优秀 Galgame 一样的存在。但是「MyGO」的二创给我带来了各种全新的体验。比如我开始了解角色的声优,了解所谓的 2.5 次元偶像企划。为了现场见声优,我甚至年底久违地去了次漫展。本来几乎不唱歌的我,也开始努力把「MyGO」的歌曲唱得好听。这些新的体验对我来说也仅仅是调剂,我不会像某些过分热情的粉丝一样花太多精力在这上面,但我很高兴能够认识这些全新的体验,去看到生活中我不曾见过的风景。

看过「MyGO」后我又自然而然地看了「Grils Band Cry」(GBC)。GBC 也是讲少女乐队的故事,但风格更偏现实,更沉重一些。剧中的人物,一直在强调一个观点:要创作自己的心声,而非迎合观众。这同样也是我一直以来的信念。这个世界不断在教人们做「正确」的事情,强迫人们放其自己本来想做的事情。不知道是幸运还是不幸,我从来就没有从迎合世界中获取过恩惠。等回过神来的时候,我的心中只剩下了对世界的复仇之情。我不想被任何世俗的负担左右,我只想做我想做的事情,我想证明我是对的。我想做的事情并不是真正意义上从我自己出发,只会有利于我的事情。就像 GBC 中的女主角被一首歌鼓舞了一样,我也是在种种鼓舞中走过来的人。我想做的,是再现这些对他人的鼓舞。所以,我所追求的,并不是只有我在追求,是很多人想追求,却因种种原因难以伸手把握的事物。既然如此,我会从我开始,斩断这些世俗的枷锁。

我今年看的这两部动漫助长了追梦的热情。我对它们的感谢无以言表,所以只能用实际行动来表达。就像现在写这篇文章的时候一样。我不想思考是不是太晚了要睡觉,不去害怕会不会写不完。只要心中还有一丝冲动,我就会继续诉说,直到心满意足了为止。


最后聊今年玩的视觉小说类游戏。由于今年「逆转裁判 4 5 6 」上架了 Steam,所以我把之前没能在模拟器上玩的《逆转裁判5》玩完了,还顺便玩完了《大逆转裁判》两部曲。感叹于巧舟的才华后,我又玩了他的《幽灵诡计》。这些作品质量都还不错,但由于我早就玩过前几部「逆转」作品,那股新鲜感已经没了。我还顺便玩了之前剩下的《弹丸论破V3》,只能说作者有创作故事的才华,却完全不会做游戏。

年底赶论文之余,我竟然玩完了老版「寒蝉」、老版《月姬》、《魔法使之夜》。「寒蝉」写了超长的日常,却不会让人感觉无趣(像 Muv Luv 本体那样)。整体故事恐怖悬疑氛围渲染到位,故事框架及伏笔回收还算合理。对我而言,故事中并没有太多感人的表达,但有不少引人思考的表达。我非常认可它的质量,能够理解为什么它作为画风简陋的同人作品能受到那么多人的喜爱。相比之下,《月姬》画风稍好,虽然故事偏短,但每条结局都做到了叙事完整。由于我是先玩过「FSN」的,玩《月姬》时能够在画风和叙事上体会到熟悉感,这种体验很有趣。正如很多其他评论一样,相比「FSN」,《月姬》有一种清冷感,那淡淡的氛围更人回味无穷。至此,时隔二十年,我总算欣赏完了当年的「同人三大奇迹」。好像这个时间有点晚。但可能时间的早晚并不是很重要,因为那份同(独)人(立)创(游)作(戏)的创作理念已经传达到了我这里。我会将它们传递下去。《魔法使之夜》就是成熟的商业作品了,一流的美术和演出,笔法成熟的作家,按理来说会给人舒适的体验。但在我看来,这部作品显然是未完成之作,很多东西没讲完故事就戛然而止了。我不认为未完成的作品是好作品。

我一般通过玩游戏来休闲。但仔细想来,还是看视觉小说时那种沉醉的心情更令我放松。我其实是喜欢文字的,但一直没有胆量去探索,只能局限于和游戏略有交集的 Galgame。我希望从明年起能够更深入地挖掘出自己的喜好,去欣赏更多优秀文学作品。


科研就没什么可以讲的了。科研是工作,并不让人讨厌,只能说是空气一般自然的存在。既然是空气,也没有太多可以讨论的必要了。倒是今年写论文时的感想我过段时间会正式发表出来。

关于科研,我最想说的是:哪怕科研是工作,是一份非常有趣的工作,我也不想让它不必要地占用我做其他事情的时间。无论会承担怎样的后果,我都要花时间去做其他想做的事情。


按照习惯,我还是收个尾吧。我不想花太多的理性去思考,为什么要写这篇文章,这篇文章是写给谁看,到底写完了有什么意义。只是昨天睡着之前想起这几天看了一些别人的年终总结,突发奇想要把自己今年的事情也整理一下而已。于其花几个小时去思考做这件事的意义,不如像现在这样就动笔写完了。不需要有什么意义,仅凭一时的冲动就够了。是的,不经深思熟虑的冲动会导致失败。但是,我心中总有一些坎是不靠冲动就跨不过去的。如果说深思熟虑后,仍然愿意面对困难的事物可以称之为勇敢的话,那么先不顾一切地做出行动,再为失败负责,同样也是勇敢的。可能明早开始我也会羞得不敢看这篇文章,那也没什么关系。我不会许诺说明年年底还会写这样的文章,我不想给自己添加任何负担。想到做就做,想做就做,不想做就不做,我要尽力地去自由地活下去。

由于我严谨的作风,我又把文章通读了一遍,忍不住站在读者视角重新感受了一下这篇文章。说实话,我文字功底太差了,没有足够的写作经验,不能把复杂深刻的情感充分表达。哪怕我写的时候自己的情感已经酝酿得很到位了,读起来还是太理性了。但没什么关系,就凭我的洞察力和自我进化能力,提升文字表达能力是分分钟的事情。是啊,文字的表达能力太有限了,不管怎么描写,都只能勾勒出复杂思绪的一角。在那表层的文字之下,究竟藏住了什么呢?不付出大量的心思,是猜不到的。也因此,无法期盼读者能够有耐心猜下去。那么,就不断地强化笔锋的锐度,不断地写下去吧。

今年年初,多尺度自回归模型 VAR 为图像生成开辟了新的发展方向:通过将图像生成建模成下一尺度预测,且每轮一次性生成同一尺度的所有像素,VAR 以极快的速度实现了高质量图像生成。随后,有许多工作都尝试对其改进。为弥补 VAR 中 VQ (Vector Quantization,向量量化) 操作引入的信息损失,HART (Hybrid Autoregressive Transformer,混合自回归 Transformer) 把 VQ 损失的信息用一张残差图表示,并用一个轻量的扩散模型来生成该残差图。做完改进后,作者用 HART 实现了 $1024 \times 1024$ 高分辨率文生图任务。在这篇博文中,我们将学习 HART 的核心方法并分析它在文生图任务上的实验结果。

论文链接:https://arxiv.org/abs/2410.10812

以往工作

本文涉及的所有自回归图像生成方法都起源于 VQVAE, VQGAN。在阅读本文前,建议读者先熟悉这两个经典工作。

HART 直接基于 VAR (Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction) 开发,且其部分思想和 MAR (Masked Autoregressive models,出自论文 Autoregressive Image Generation without Vector Quantization) 类似。欢迎大家阅读我之前的解读。

VAR 解读

MAR 解读

在 VQGAN 两阶段生成方法的基础上,VAR 让自编码器输出一系列不同尺度的图像词元 (token),而不仅仅是最大尺度的词元。生成时,VAR 自回归地生成不同尺度的词元图,同一尺度的词元图会在一轮 Transformer 推理中一次性生成。

VQ 操作会丢失编码器输出中的信息,这导致所有使用 VQ 自编码器的图像生成模型生成质量略低。VAR, VQGAN 等方法之所以不得不使用 VQ,是因为这些方法都用类别分布(categorical distribution)来建模词元的分布。为了彻底去除 VQ 操作,MAR 使用扩散模型来代替类别分布,使得我们能够用精度更高的 VAE 来压缩图像。

弥补 VQ 的信息损失

为了缓解 VAR 中 VQ 造成的质量下降,HART 使用了一项思路直接的设计:既然 VQ 无论如何都会造成信息损失,不妨把损失的信息看成一张残差图像。用普通的 VAR 生成完图片后,再用扩散模型生成该残差图像。把残差图像加到原输出图像后,新输出图像会质量更高。

让我们通过论文里的图片来直观感受这一点。第一行是 VAR 自编码器和 HART 的混合自编码器的重建结果。可以看出,由于 VQ 操作,模型难以重建输入图像。第二行原 VAR 的输出和残差图像输出。我们发现,加上残差图像后,图像的细节更加丰富,不会像之前一样模糊。

在下两个小节里,我们来学习 HART 是怎么分别改进 VAR 的词元生成模型和自编码器的。

用扩散模型生成残差图像

为了理解整套方法,我们需要理解 HART 的「残差图像」是从哪来的。因此,我们先看词元生成模型上的修改,再看自编码器的对应修改。

我们先仔细回顾一下 VAR 中 VQ 误差是怎么引入的。VAR 借用了传统拉普拉斯金字塔的思想来建模不同尺度的词元图。

也就是说,VAR 并没有将完整图像拆解成内容相同、不同分辨率的词元图,而是拆解成了最低分辨率的图以及各个尺度上的信息损失。这里的信息损失不仅包括了下采样导致的,还包括了 VQ 导致的。

即使在多尺度拆解时考虑了 VQ 的信息损失,最终的重建特征(即解码器输入,词元查表输出的累加)依然不能和编码器输出特征完全一致。HART 想用扩散模型生成的「残差图像」,就是上图中重建特征和编码器输出特征的差。

和离散的词元图不同,残差图像是连续的。为了生成该连续图像,HART 参考 MAR,使用了一个图像约束的扩散模型。该任务可以解释为:已知离散词元图的输出,该如何用扩散模型生成细节,以提升输出图像质量。

HART 的生成模型示意图如下所示。前面的生成过程和 VAR 一模一样。在最后一步,Transformer 的中间隐状态会输入给用 MLP 表示的扩散模型,扩散模型会为每个词元独立地预测残差量。也就是说,这不是一个图像扩散模型,而是只生成一个词元值的像素扩散模型。词元之间的采样互相独立。得益于这种独立性假设,HART 可以用一个非常轻量的扩散模型来生成残差图,几乎没有增加整体的生成时间。

HART 还将 VAR 的类别约束换成了文本约束。我们稍后在实验部分讨论。

AE + VQVAE 混合自编码器

知道了 HART 要生成的残差图像从何而来,我们可以回头学习自编码器上的对应修改。现在,自编码器的解码器有两种输入:一种是 VAR 离散词元累加而成的近似重建特征,一种是加上了 HART 的残差图的精确重建特征,这个重建特征就等于编码器输出特征。为了同时处理这两类输入,在训练 HART 的混合自编码器时,解码器的输入一半的时候是编码器输出,另一半的时候是离散词元的重建特征。当然,在生成时,由于加上了残差图像,可以认为解码器的输入就等于编码器的输出。

下图中采用的术语 token 与 VAR 不同。VAR 把编码器输出和解码器输入都叫做特征图 (feature map),把过了 VQ 操作的索引图叫做词元图 (token map)。而 HART 将 VAR 里的特征图称为 token,continuous token 表示编码器输出特征,discrete token 表示词元的重建特征。这篇博文采用了 VAR 的称呼方法。同理,HART 里的 residual token 在本文被称为「残差图像」。

这样看来,HART 的混合编码器既像没有 KL Loss 的 VAE,即普通自编码器 (AE),也像 VQVAE。

高分辨率文生图实现细节

我们来简单看一下 HART 是如何把 ImageNet $256 \times 256$ 按类别生成的 VAR 拓展成 $1024 \times 1024$ 的文生图模型的。

  • 文本约束:HART 没有通过交叉注意力输入文本信息,而是和 VAR 对类别嵌入的做法一样,将文本嵌入作为第一尺度的输入及 AdaLN 层的输入。
  • 位置编码:不管是对于尺度编号还是图像位置编号,VAR 用的是可学习的绝对位置编码。HART 对尺度采取了正弦编码,对图像词元采取了 2D RoPE(旋转位置编码)。
  • 更大尺度:原 VAR 词元图的最大边长是 16,HART 往后面添加了 21,27,36,48,64 这几个边长。
  • 轻量级扩散模型:由于扩散模型仅需建模单个词元的分布,它仅有 37M 参数,只需 8 步就能完成高质量采样。

定量实验结果

先看一下最热门的「刷点」指标——ImageNet $256 \times 256$ 按类别生成。作者没放最好的 MAR 模型,我补上去了。

在这个任务上,HART 和 VAR 的主要区别在于是否使用扩散模型输出残差图像。从结果可以看出,残差扩散模型几乎没有提升推理时间,却对 FID 指标有不小的提升(考虑到数值越低,提升难度越大)。并且,通过比较不同模型的速度,我们发现类 VAR 模型最大的优势在于推理速度快。

再看一下这篇论文重点关注的文生图指标。除了常用的主要衡量文图匹配度的 GenEval 外,论文还展示了两个今年刚出的指标: MJHQ-30K 数据集上的指标和 DPG-Bench。

这些指标不见得很有说服力。在由用户投票的排名中 (https://imgsys.org/rankings),Playground v2.5 最好,SD3 和 PixelArt-Σ 差不多。但是,MJHQ FID 和 DPG-banech 指标都不能反映出这些模型的排名。特别地,FID 用到的 Inception V3 网络是在 $299 \times 299$ 的 ImageNet 上训练的,所以 FID 既不能很好地反映高分辨率图像的相似度,也不能反映更复杂的图像的相似度。

综上,HART 在高分辨率文生图任务上的表现暂时不能通过实验结果反映。根据部分社区用户的反馈(https://www.reddit.com/r/StableDiffusion/comments/1glig4u/mits_hart_fast_texttoimage_model_you_need_to_see/ ),HART 在高频细节的生成上存在缺陷。通过回顾 HART 的方法,我们可以猜测这是残差扩散模型的设计不够好导致的。

总结

为了缓解 VQ 自编码器中 VQ 操作带来的信息损失,HART 把信息损失当成一张残差图,并额外用一个轻量级像素扩散模型来独立地生成残差图的每个像素。HART 把这一改进直接应用到了 VAR 上,并提升了 VAR 的 ImageNet FID 指标。HART 在高分辨率文生图任务上依然无法媲美扩散模型,并且由于扩散模型存在诸多加速手段,它在生成速度上也没有优势。

VQ 操作将复杂的图像转换成了易于学习的图像词元,但牺牲了自编码器的重建质量。为了改进这一点,有许多工作都试图改进原 VQVAE 的最近邻 VQ 操作。但无论如何,VQ 导致的误差是不可避免的。HART 从另一个角度缓解 VQ 重建误差:用另一个模型来生成残差图像。这种设计思想很有前途,有希望彻底去除 VQ 的误差。然而,天下没有免费的午餐,提升了生成效果,就不得不增加训练和生成时间。HART 用轻量级像素扩散模型生成残差图的做法虽然不会拖慢模型速度,但效果还不够好。或许可以将其换成一个感受野稍大一点的扩散模型,在不显著增加生成时间的前提下提升残差图生成效果。

今年四月,北大和字节跳动在 Arxiv 上发表了论文 Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction,介绍了一种叫做 Visual Autoregressive Modeling (视觉自回归建模,VAR)的全新图像生成范式。这种自回归生成方法将高清图像用多尺度词元图像表示,并用下一尺度预测代替了此前常用的下一词元预测。在 ImageNet $256 \times 256$ 图像生成任务上,VAR 的表现超越了 DiT。我们组的同学第一时间看了这篇论文,大家都觉得这篇论文有不小的创新,但其方法能否完全代替扩散模型还有待验证。通常来说,这篇论文的关注度会逐渐降下去,但近期发生的两件大事将 VAR 论文的热度推向了空前的高度:论文一作的严重违纪行为招致字节跳动对其索赔 800 万元、论文被评选为 Neurips 2024 会议的最佳论文。借此机会,我决定认真研究一下这篇论文并把我的学习结果分享给大家。

在这篇博文中,我会先回顾与 VAR 密切相关的早期工作 VQVAE 和 VQGAN,再介绍论文的方法细节与实验结果,最后分享我对该工作的测试结果与原理探究。在读 VAR 论文时,我发现有个地方的设计存在缺陷。相关实验结果表明, VAR 论文并没有完整地分析出这套方法有效的原因。欢迎大家仔细阅读这一部分并提出自己的思考与见解。

论文链接:https://arxiv.org/abs/2404.02905

VQGAN 原理回顾

VAR 算是 VQGAN 工作的改进版,而 VQGAN 又是 VQVAE 工作的改进版。要了解 VAR 的背景知识,最直接的方法就是回顾 VQVAE 与 VQGAN 这两个经典工作。我们先从自回归这种生成范式开始聊起,再将目光移向图像自回归生成,最后复习 VQVAE, VQGAN, Transformer 的实现细节。

推荐大家阅读我之前写的相关博文:

轻松理解 VQ-VAE:首个提出 codebook 机制的生成模型

VQGAN 论文与源码解读:前Diffusion时代的高清图像生成模型

Attention Is All You Need (Transformer) 论文精读

图像自回归生成

自回归(Autoregressive)是一种直观易懂的序列生成范式:给定序列前 $n$ 个元素,模型输出第 $n+1$ 个元素;把新元素添加进输入序列,再次输出第 $n+2$ 个元素……。以下是文本自回归生成的一个示例:

1
2
3
4
(空)  -> 今
今 -> 天
今天 -> 早
今天早 -> 上

具体来说,模型的输出并不是下一个元素应该是什么,而是下一个元素可能是什么。也就是说,模型的输出是下一个元素的概率分布。通过不断对下一个元素采样,我们就能随机生成出丰富多样的句子。

自回归生成仅适用于有顺序的序列数据。为了用自回归生成图像,我们需要做两件事:1)把图像拆分成一个个元素;2)给各个元素标上先后顺序。为此,最简单的做法是将图像拆成像素,并从左到右,从上到下地给图像生成像素。比如下图是经典自回归图像生成模型 PixelCNN 的示意图。假设图像有 $3 \times 3$ 个像素,并按顺序从左上到右下标号。在生成第 $5$ 个像素时,模型只能利用已经生成好的前 $4$ 个像素的信息。模型的输出是一个概率分布,表示灰度值大小分别取 0, 1, ..., 255 的概率。

顺带一提,建模概率分布的方法有很多种,这里我们使用的分布被称为类别分布(categorical distribution)。这种方法的好处是形式简洁,可以用简单的算法采样,缺点是元素的取值必须是离散的。比如虽然图像的灰度值理论上可以取 0~1 中间的任何实数(假设灰度值被归一化了),但我们用上图所示的 PixelCNN 时,只能表示 0, 1/255, 2/255, ..., 1 这 256 种灰度值,而不能表示更加精确的值。

VQVAE

PixelCNN 虽然能做图像生成,但它的效率太慢了:由于像素是逐个生成的,要生成几个像素,就要运行几次神经网络。能不能加速生成过程呢?如果要生成的图像更小一点就好了。

为了加速 PixelCNN,借助图像压缩网络,VQVAE 工作提出了一种两阶段的图像生成方法:先生成压缩图像,再用图像压缩网络将其复原成真实图像。由于压缩图像的像素数较少,而复原压缩图像的速度又很快,整套生成方法的速度快了很多。

以下是一个 VQVAE 的生成示例。根据 PixelCNN 输出的类别分布,我们可以采样出一些由离散值构成的压缩图像。这些离散值就和 NLP 里的文字一样,每一种值都有一种特殊的含义。我们可以认为离散值表示原始图像中一大块像素的颜色。借助图像压缩网络的解码器,我们可以把压缩图像复原成清晰的原始图像。

VQVAE 的训练顺序和生成顺序相反。我们先训练一个图像压缩网络。这种由编码器和解码器组成的图像压缩网络被称为自编码器,压缩出来的图像被称为隐图像(latent image)。训练好了自编码器后,我们再把训练集的所有图像都转成隐图像,让 PixelCNN 学习生成隐图像。比较有趣的是,训练 PixelCNN 时,只会用到编码器;而生成时,只会用到解码器。

在上述讨论中,我们略过了一个实现细节:该怎么让网络以离散值为输入或输出呢?输入离散值倒还好办,在 NLP 中,我们用嵌入层把离散的词语变成连续向量,这里的做法同理。可怎么让网络输出离散值呢?这里就要用到向量离散化(vector quantization, VQ)操作了。

离散化操作我们都很熟悉,将小数四舍五入至整数就是一种最常见的离散化。四舍五入,本质上是让一个小数变成最近的整数。同理,对于向量而言,假设我们已经准备好了一些向量(对应前面的「整数」),那么向量离散化就表示把输入的任意向量变成最近的已知向量。这里的「最近」指的是欧几里得距离。

具体示例如下所示。编码器可以输出一个由任意向量构成的二维特征。通过查找嵌入层里的最近邻,这些任意的向量会被转换成整数,表示最近邻的索引。索引可以被认为是 NLP 里的词元 (token),这样编码器输出特征就被转换成了词元构成的隐图像。而在将隐图像输入进解码器时,我们把嵌入层当成一张表格,利用隐图像里的索引,以查表的形式将隐图像转换成由嵌入构成的特征。准确来说,这个把图像压缩成离散隐图像的自编码器才被叫做 “VQVAE”,但有时我们也会用 VQVAE 代表整套两阶段生成方法。

上图中的「编码器输出特征」、「词元」、「嵌入」在不同论文里有不同的叫法,且一般作者都只会用数学符号来称呼它们。这里我们用了 VAR 论文的叫法。

嵌入层的具体学习过程我们不在此展开,对这块知识不熟悉的读者可以去仔细学习 VQVAE 论文。

VQGAN

VQVAE 的效果并不理想,这是因为它的压缩网络和生成网络都不够强大。为此,VQGAN 工作同时改进了 VQVAE 的两个网络。

  • VQGAN 工作将离散自编码器 VQVAE 换成了 VQGAN。在 VQVAE 的基础上,VQGAN 在训练时添加了感知误差和 GAN 误差,极大提升了自编码器的重建效果。
  • VQGAN 工作还把生成模型从 PixelCNN 换成了 Transformer。

Transformer

Transformer 是目前最主流的主干网络。相比其他网络架构,Transformer 的最大特点是序列里的元素仅通过注意力操作进行信息交互。因此,为了兼容文本自回归生成任务,最早的 Transformer 使用了两个特殊设计:

  • 由于注意力操作不能反映输入元素的顺序,词元嵌入在输入进网络之前,会和蕴含了位置信息的位置编码相加。
  • 自回归生成要求之前的词元不能看到之后的词元的信息。为了控制词元间的信息传播,Transformer 给自注意力操作加上了掩码。

VQGAN 用了完全相同的设计,把图像词元当成文本词元用 Transformer 来生成。

从词元预测到尺度预测

上述的传统图像自回归生成都是采用下一个词元预测策略:

  • 将图像用自编码器拆成离散词元。
  • 从左到右、从上到下按顺序逐个生成词元。

尽管通过自编码器的压缩,要生成的词元数已经大大减少,但一个个去生成词元还是太慢了。为了改进这一点,VAR 提出了一种更快且更符合直觉的自回归生成策略:

  • 将图像用自编码器拆成多尺度的离散词元。比如,原来一张隐图像的大小是 $16 \times 16$,现在我们用一系列尺度为 $1 \times 1, 2 \times 2, …, 16 \times 16$ 的由词元构成的图像来表示一张隐图像。
  • 从最小的词元图像开始,从小到大按尺度生成词元图像。

在这种策略下,我们要同时修改自编码器和生成模型。我们来看一下 VAR 是怎么做的。

多尺度残差离散自编码器

先来看自编码的修改。现在词元图像不是一张图像,而是多张不同尺度的图像。由于词元图像的定义发生了改变,编码器特征和嵌入的定义也要发生改变,如下图所示。

向量离散化部分我们可以沿用 VQVAE 的做法。现在新的问题来了:编码器的输出和解码器的输入都只是一张图像。该怎么把多尺度的图像组合成一张图像呢?

最简单的做法是完全不修改编码器和解码器,还是让它们输入输出最大尺度的图片。只有在中间的向量离散化/查表部分,我们才把这些图片下采样。

VAR 用了一种更加高级的做法:用残差金字塔来表示这些隐空间特征。我们先来回顾一下拉普拉斯金字塔这一经典图像处理算法。我们知道,图像每次下采样的时候,都会损失一些信息。既然如此,我们可以将一张高分辨率的图像表示为一张低分辨率的图像及其在各个分辨率下采样后的信息损失。如下图所示,最右侧的一列表示拉普拉斯金字塔的输出。

在计算拉普拉斯金字塔时,我们不断下采样图像,并计算当前尺度的图像和下一尺度的复原图像(通过上采样复原)的残差。这样,通过不断上采样最低尺度的图像并加上每一层的残差,我们最终就能精准复原出高分辨率的原图像。

现在,我们想把类似的金字塔算法应用到编码器特征上。该怎么把最大尺度的编码器特征拆解成不同尺度的图像的累加呢?

在计算拉普拉斯金字塔时,本质上我们用到了两类操作:退化和复原。对于图像而言,退化就是下采样,复原就是上采样。那么,对于编码器输出的隐空间特征,我们也需要定义类似的退化和复原操作。比较巧妙的是,VAR 并没有简单地把退化和复原定义为下采样和上采样,而是参考 Autoregressive Image Generation using Residual Quantization 这篇论文,将向量离散化引入的误差也算入金字塔算法的退化内。也就是说,我们现在的目标不是让编码器特征金字塔的累加和编码器特征相等,而是想办法让嵌入金字塔的累加和编码器特征尽可能相似,如下图所示。

基于这一目标,我们可以把退化定义为下采样加上离散化、查表,复原定义成上采样加一个可学习的卷积。我们来看看在这种新定义下,原来 VQVAE 的向量离散化操作和查表操作应该怎么做。

先看新的多尺度向量离散化操作。这个操作的输入是编码器特征,输出是一系列多尺度词元图像。算法从最低尺度开始执行,每个循环输出当前尺度的词元图像,并将残差特征作为下一个循环的输入特征。

对于多尺度查表操作,输入是多尺度词元图像,输出是一张最大尺度的隐空间特征,它将成为自编码器的解码器的输入。在这步操作中,我们只需要分别对各个尺度的词元图像做查表和复原(上采样+卷积),再把各尺度的输出加起来,就能得到一个和编码器特征差不多的特征。注意,为了方便理解,这几张示意图都省略了部分实现细节,且一些数值不是十分严谨。比如在查表时,我们可以让不同尺度的词元共享一个嵌入层,也可以分别指定嵌入层。

总结一下这一小节。为了实现尺度自回归生成,我们需要把图像编码成多尺度的词元图像。VAR 采用了一种多尺度残差离散化操作:将编码器特征拆解成最小尺度的特征以及不同尺度的残差特征,并对不同尺度的特征分别做向量离散化。这种做法不仅能高效地将特征拆解成多个尺度,还有一个额外的好处:原来 VQVAE 仅对最大尺度的特征做向量离散化,离散化后的误差会很大;而 VAR 把向量离散化引入的误差分散到多尺度离散化中,巧妙地降低了离散化的误差,提升了 VQVAE 的重建精度。

下一尺度自回归生成

把图像压缩成多尺度词元图像后,剩下的事就很简单了。我们只需要把所有词元拆开,拼成一维词元序列,之后用 Transformer 在这样的序列上训练即可。由于现在模型的任务是下一尺度预测,模型会一次性输出同尺度各词元的概率分布,而不是仅仅输出下一个词元的。这样,尽管序列总长度变长了,模型的整体生成速度还是比以前快。同时,随着预测目标的变更,自注意力的掩码也变了。现在同尺度的词元之间可以互相交换信息,只是前一尺度的词元看不到后面的词元。以下是一个 $3 \times 3$ 词元图像在下一词元和下一尺度预测任务下的注意力掩码示意图及生成过程示意图。

除此之外,VAR 的 Transformer 还做了一些其他修改:1)除了给每个词元加上一维位置编码外,同一尺度的词元还会加上同一个表示尺度序号的位置编码。所有位置编码都是可学习的,而不是预定义的正弦位置编码。2)Transformer 与解码器的共用嵌入层。另外,在生成新一层时,为了复用已经生成好的图像的信息,新一层的初始嵌入是通过对上一层的生成结果 bicubic 上采样得到的。

该 Transformer 的其他设计都与 VQGAN 相同。比如,Transformer 采用了 decoder-only 的结构。为了添加 ImageNet 类别约束,第一层的输入是一个表示类别的特殊词元。训练时用的误差函数是交叉熵函数。

ImageNet 图像生成定量实验

VAR 的方法部分我们看得差不多了,现在来简单看一下实验部分。论文宣称 VAR 在图像生成实验和参数扩增实验上都取得了不错的成果。特别地,VAR 的拟合能力胜过了 DiT,生成速度是 DiT 的 45 倍以上。我们就主要看一下 VAR 在 ImageNet $256 \times 256$ 图像生成上的实验结果。以下是论文中的表格。我同时还附上了何恺明团队的 MAR 工作(Autoregressive Image Generation without Vector Quantization)的实验结果。

先比一下 DiT 和 VAR。先看速度,不管是多大的模型,DiT 的速度都远远慢于 VAR。再看以 FID 为代表的图像拟合指标。VAR 在参数量为 600M 左右时并没有 DiT 效果好。但继续增加参数量后,DiT 的 FID 没有变好的趋势,而 VAR 的 FID 一直在降。最终 VAR 的 FID 甚至超过了 ImageNet 的验证集,可以认为 FID 再低的也意义不大了。

再比一下 MAR 和 VAR。MAR 的刷指标能力更加恐怖,943M 的模型就能有 1.55 的 FID。但根据 MAR 论文,其速度是 DiT-XL 的 5 倍左右,也就是说 VAR 还是比 MAR 快,是 MAR 速度的 9 倍左右。

ImageNet 图像生成已经被各个模型刷到头了。FID 结果能说明 VAR 的拟合能力很强,最起码不逊于 DiT。但在更有挑战性的文生图任务上,VAR 的效果还有待验证。另外,虽然刷指标的时候 DiT 用了 250 步采样,但实际用起来的时候一般就是采样 20 步。如果算上蒸馏的话,采样步数能缩小到 4 步。加上这些加速技巧的话,VAR 不见得会比 DiT 快。

VAR 各尺度生成结果

看完了论文的主要内容,我来分享一下我对 VAR 的一些理论分析与实验结果。

先看一下随机采样结果。我用的是最大的 d=30 的 VAR 模型。在官方采样脚本的默认配置下,两个随机种子 (0, 15) 的输出如下所示。用到的图像类别为火山、灯塔、老鹰、喷泉,每个类别的图各生成了两张。图像的生成速度很快,一秒就生成了全部 8 张图片。

我们还可以观察每个尺度的生成结束后解码出的临时图片。和我们预估得一样,图像是按从粗到精的顺序逐渐生成的。

为了进一步探究每一个尺度负责生成哪些图像成分,我们可以做如下的实验:从某个尺度开始,随机更换新的随机数生成器。这样,每张动图里不变的部分就是前几个尺度生成好的内容;不断在变的部分就是后几个尺度负责的内容。可以看出,从第三个尺度开始,图像的内容就基本固定下来了,也就是说结构信息是在前两个尺度里生成的。越往后,图像的细节越固定。

这个结果还挺令人惊讶的:难道 $2 \times 2$ 这么小的特征图就已经决定了图像的整体内容?让我们来仔细探究这一点。

有缺陷的单尺度生成

不知道大家在学习 VAR 的采样算法时候有没有感到不对劲:在生成同一个尺度的词元图像时,每个词元是独立地在一个概率分布里采样。

而根据作者在论文里的说法,VAR 的尺度自回归是一种新的自回归概率模型:

其中,$r_k$ 表示从小到大第 $k$ 个尺度的词元图像,共 $K$ 个尺度。同一个尺度的词元图像 $r_k$ 的每个词元的分布是并行生成的。这也就是说,VAR 的这种训练(用交叉熵误差)和采样方式是认为每张词元图像的概率等于所有词元的概率的乘积,词元的分布之间是独立的:

其中,$r_k^i$ 表示第 $k$ 个尺度的第 $i$ 个词元,$I_k$ 为第 $k$ 个尺度的词元总数。我觉得上面这个等式是不成立的,哪怕有之前尺度的信息作为约束,同一尺度的每个词元的概率分布之间不会是互不相关的。且随着 $I_k$ 的增大,上面这个式子的误差会越来越大。

词元之间的采样互相独立,理论上会导致图像出现不连贯的地方。比如,假设一个图像词元表示 $16 \times 16$ 个像素,那么每隔 16 个像素图像就会出现「断层」。但是,为什么 VAR 的输出结果那么正常呢?仔细分析 VAR 的生成算法,我们可以发现有两项设计提升了图像的连续性:

  • VAR 的自编码器使用了向量离散化操作。这个操作会让解码器的输入总是合理的,解码器也总是会输出连贯的图像。
  • 在生成一个新尺度的图像时,模型输入会被初始化成之前尺度的图像的 bicubic 上采样。bicubic 采样保证了词元嵌入之间的连续性。

此外,为了进一步缓解独立采样带来的负面影响,VAR 在生成完第二或第三个尺度后就已经把图像的整体内容确定下来了,后面的生成只是略微影响图像细节而已(因为随着词元数量变多,独立采样的误差越大)。这个结论已经在前文的可视化结果中验证了。为了证明只有前几个尺度是重要的,我做了一个大胆的实验:用 Transformer 生成完前两个尺度的词元后,后续所有词元都随机生成。如下图所示,我展示了固定前两个尺度的输出后,多个随机种子下的生成结果。结果显示,如果前两个尺度的词元生成得比较好,后面词元无论采样得多乱,都不怎么会影响最终的图像质量。

根据这些实验结果,我认为 VAR 真正有效的原因并不能用「下一尺度预测这种全新生成范式更好」这样粗浅的话来概括。VAR 中最核心的组件可能是其多尺度残差离散自编码器。这个编码器至少做到了以下几件事:

  • 使用向量离散化确保解码器的输入总是合理的。
  • 使用多尺度残差设计,且下一尺度的残差图像不仅记录了因下采样而导致的信息损失,还记录了因向量离散化带来的精度损失。相比简单的、人类能够理解的拉普拉斯金字塔,这种可学习的多尺度拆分方法或许更加合理。
  • 使用 bicubic 对低尺度词元图上采样。这步固定的操作让生成的图像总是连续的。

当然,这几件事是互相耦合的。不进行更深入的实验的话,我们难以解耦出 VAR 中最有效的设计。

多尺度生成其实并不是什么新奇的思想。之前 StyleGAN 和 Cascaded Diffusion 都用了类似的策略。然而,VAR 做了一个大胆的设计:同一尺度的不同词元在采样时是相互独立的。令人惊讶的是,这种在数学上不太合理的设计没怎么降低图像的质量。并且,得益于这一设计,VAR 能够并行地对同一尺度的词元采样,极大地提升了生成速度。

总结与评论

此前,以经典工作 VQGAN 为代表的图像自回归生成模型无论在速度上还是图像质量上都不尽如人意。究其原因,下一个图像词元预测的建模方式既不够合理,也拖慢了生成速度。为此,VAR 提出一种新式自回归策略:将词元图像拆分成多个尺度,通过下一尺度预测实现图像生成。为了兼容这一设计,VAR 对 VQGAN 的自编码器和 Transformer 都进行了修改:自编码器能够将图像编码成多尺度的残差词元图像,而 Transformer 同时输出同一尺度每个词元的独立分布。实验表明,VAR 在 ImageNet 图像生成指标上超越了以 DiT 为代表的扩散模型,且生成速度至少比 DiT 快 45 倍。另外,还有实验表明 VAR 符合扩增定律:增加参数量即可提升模型性能。

我个人认为,和其他前沿生成模型一样,VAR 在 ImageNet 上的表现已经满分了。它能否完成更困难的图像生成认为还有待验证。最近字节跳动发布了 VAR 的文生图版本:Infinity,但这个模型还没有开源。我们可以持续关注 VAR 的各个后续工作。VAR 的生成速度也没有比 DiT 快上那么多,通过减小采样步数,再加上模型蒸馏,DiT 不会比 VAR 慢。当然,VAR 或许也存在进一步加速的可能,只是相关研究暂时没有扩散模型那么多。

VAR 的数学模型是存在缺陷的:词元图的分布不应该等于词元间的独立分布的乘积。最起码论文里没有任何相关分析(用了类似做法的 MAR 论文也没有分析)。通过一些简单的生成实验,我们发现由于 VAR 在其他设计上提升了输出图像的连续性,哪怕同一尺度的词元间是独立采样,甚至是随机均匀采样,模型的输出质量也不会太差。我们需要通过更深入的实验来挖掘 VAR 的生效原理。

我觉得如果一个科研工作能够解释清楚 VAR 中哪些模块起到了最主要的作用,并取其精华,去其糟粕,提出一个更好的生成模型,那这会是一个很不错的工作。我觉得能够探索的方向有:

  • VAR 的前几个尺度的词元图是最重要的。能不能用更好的方式,比如用扩散模型,来生成前几个尺度的图像,而更大尺度的词元图用一个比 Transformer 更高效的模型来生成。这样模型的质量和效率能进一步提升。
  • VAR 还是用了 VQ 自编码器。无论怎么样,VQ 操作都会降低模型的重建质量。但另一方面,VQ 也能起到规范解码器输入的作用。究竟我们能不能把 VQ 自编码器换成精度更高的 VAE 呢?换了之后怎么设计多尺度编码呢?

以前我一直将 “quantization” 翻译成「离散化」,但经仔细学习后发现它与真正的离散化 “discretization” 存在区别。因此,以后我会采用「量化」这个更常见的翻译,尽管我认为这个翻译容易和「量化投资」里的「量化」混淆。

按分辨率从低到高的顺序生成图像是一种常见思路。此外,Diffusion Forcing 等论文带来了一种新的扩散模型视频生成思路:将视频生成转换为约束于前几帧图像的单张图像自回归生成。Pyramid Flow 工作把两种思路结合起来,提出了一种新的视频生成范式:在自回归生成新帧时,用低分辨率的前几帧图像作为约束。这使得模型能够更高效率地利用历史帧的信息。

论文名:Pyramidal Flow Matching for Efficient Video Generative Modeling

Arxiv: https://arxiv.org/abs/2410.05954

GitHub:https://github.com/jy0205/Pyramid-Flow

以往工作

将图像生成拆解成从低分辨率到高分辨率是一种很常见的思想。基于扩散模型,有多种方式来应用这种思想。一种比较直接的方式是显式将图像生成分解成生成最低分辨率的图像和多轮超分辨率,代表工作是 Cascaded Diffusion Models for High Fidelity Image Generation;另一种更加巧妙的方式是将图像上采样和扩散模型的去噪同时进行,代表工作是 f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation。本文的多尺度设计和 f-DM 非常相似,我会在文末详细分析二者的区别。

将视频生成转换为约束于之前所有帧的图像生成是一种再简单不过的想法。然而,在使用扩散模型时,选择最佳的约束方式并不是一个简单的问题。比较常见的添加图像约束的方式是与原输入在通道维度上拼接。对于视频自回归生成而言,这个做法的问题是网络对于约束图像和待生成图像的处理不统一。近期,Diffusion Forcing 工作告诉我们,我们可以给视频扩散模型的不同帧添加不同的噪声,并修改注意力机制,从而将其转换成一个约束于之前帧的图像生成模型,并且模型对每一帧的处理在注意力层以外是统一的。

流匹配 (Flow Matching) 可以简单看成一种改进了噪声调度机制的 DDPM。假设时刻 $0$ 表示纯高斯噪声,时刻 $1$ 表示清晰图像,按照最常用的流匹配方法,中间 $t$ 时刻的带噪图像为二者的线性插值。

此外,去噪网络的预测目标也从残余噪声 (epsilon) 改为了速度(velocity)。

相关博文:

Diffusion Forcing

多尺度的加噪过程

扩散模型中,我们要把清晰图像中的信息逐渐破坏掉。这样的图像退化方式不只添加高斯噪声一种,我们可以在添加噪声的同时下采样图像,定义出一种新的扩散模型前向过程(退化过程)。

如下图的一维扩散模型所示,Pyramid Flow 在加噪的同时做两次下采样。总时间被拆成了三段,第 $k$ 个阶段的时间范围是 $[s_k, e_k]$。

做下采样后,图像中的信息会突然减少。为了让同一个阶段的图像逐渐失去这些信息,而不是在一次下采样中突然失去,我们采用这样的插值策略:设 $\mathbf{x}_t$ 为最大分辨率下 $t$ 时刻的带噪图像(其计算方法由上文的流匹配噪声公式决定),那么在第 $k$ 阶段的分辨率下,$t$ 时刻的带噪图像 $\hat{\mathbf{x}}_t$ 为:

其中,$Down(\mathbf{x}, a)$ 表示把 $\mathbf{x}$ 下采样 $a$ 倍,$Up(\mathbf{x})$ 表示把 $\mathbf{x}$ 上采样 $2$ 倍,$t’=(t-s_k)/(e_k-s_k)$ 表示 $t$ 时刻在第 $k$ 阶段里的插值比例。初看这些公式可能会有点头大,我们可以先看下面的示意图再回头看公式。

公式里做插值的两项分别表示当前阶段最清楚、最模糊的图像,插值比例 $t’$ 为 $t$ 在当前阶段的时间窗口里的比例。这种插值不仅让噪声强度渐变,还让因下采样而产生的信息损失渐变,就好像是我们在连续地对图像下采样一样。

修改了退化机制后,除了定义流匹配的输入变量 $\hat{\mathbf{x}}_{t}$ 外,我们还可以定义新的流匹配学习目标。流匹配的学习目标是一个速度,而速度又可以由两个端点决定。我们可以用上面的公式定义每一个阶段的端点 $\hat{\mathbf{x}}_{e_k}, \hat{\mathbf{x}}_{s_k}$,再以速度 $\hat{\mathbf{x}}_{e_k} - \hat{\mathbf{x}}_{s_k}$ 为该阶段所有点的学习目标。

重新设置上采样噪声

在多尺度生成时,我们必须仔细考虑图像应该怎么跨越两个阶段,也就是图像该怎么上采样。

大多数涉及扩散模型多尺度生成的工作(比如 Laplacian Diffusion Models sampling, f-DM)都会在上采样时考虑调整噪声强度的大小。而 Pyramid Flow 不仅考虑了噪声的强度,还考虑到了噪声的协方差矩阵——原本扩散模型不同位置的噪声是相互独立的,但 2 倍上采样后一个像素会影响到周围 4 个像素,因此噪声之间的关系需要用协方差矩阵表示而不是独立表示。作者根据协方差矩阵必须半正定这一性质推导了一个更好的修改上采样后噪声的方法。

多尺度视频生成

Diffusion Forcing 工作表明,在训练视频扩散模型时,不同帧的退化程度可以不同。而在 Pyramid Flow 的框架下,图像的退化不仅包括加噪,还包括降采样。因此,用 Pyramid Flow 训练视频扩散模型时,可以让过往帧为低分辨率的,只让最新帧为最高分辨率的。这样做可以大大减少注意力操作的输入量,提升计算效率。此外,还可以参照 Diffusion Forcing 的做法,通过因果(causal)注意力把视频扩散模型转换成约束于前几帧的单张图像自回归生成。

实现细节

最后,我们来简单看一下 Pyramid Flow 的实现细节。Pyramid Flow 可以基于任何一种 DiT 架构的文生图模型开发,比如 SD3 和 FLUX.1。为了将文生图模型转换成视频模型,本工作做了如下的适配:

  • 将图像 VAE 变成能在时间维度压缩视频的 3D VAE。
  • 在 DiT 中用 RoPE 描述时间维度的位置关系。空间维度位置编码保持原模型的设计。

另外,为了对齐不同尺度的空间位置编码,本工作参考 CogVideoX,对低分辨率图像使用了位置编码外推。

本工作开源的视频模型需要用 128 块 A100 训练约 1 天,能生成 241 帧总时长 10 秒的视频。相比其他视频模型来说要的资源已经少了很多了。

批判性分析与总结

Pyramid Flow 在多尺度加噪的设计上和以往工作 f-DM 非常相似。二者都把扩散模型的加噪拆成不同尺度,都在同一尺度内采取线性插值的方式计算带噪图片,都考虑了上采样时噪声的变化。二者的不同之处在于:

  • f-DM 用的是 DDPM 的噪声调度,而 Pyramid Flow 用的是 Rectified Flow 的公式。由于 Rectified Flow 的公式本来就是线性插值,因此它对同尺度带噪图像插值来说兼容性更好一点。
  • 基于上一点,可能是由于同尺度线性插值与 DDPM 不太兼容,f-DM 额外用一个残差项表示下采样导致的误差。Pyramid Flow 没用到这一设计。
  • f-DM 从局部信噪比的角度推导了上采样时噪声应该发生的变化,而 Pyramid Flow 是从协方差矩阵的角度。

Pyramid Flow 没引用 f-DM,作者大概是独立发明了一遍类似的加噪策略。我在 GitHub 上已经告知了作者有 f-DM 这篇类似论文。当然,哪怕是有类似工作在前,Pyramid Flow 在结合流匹配和多尺度生成上还是有一定创新的。

包括这个工作在内,有好几个工作都在用扩散模型做多尺度图像生成。但是,不像自回归中的 VAR,这些工作在纯图像生成任务上并不是很有名。因此,多尺度生成虽然是一个不错的想法,但我们还需要多多思考该怎么在扩散模型里应用它。Pyramid Flow 最大的意义可能不在于其多尺度的设计,而在于将多尺度生成融合进了视频生成中。这样做最直接的好处是减少了历史帧的数据量,提升了模型计算效率。

虽然这个工作试图用同一个去噪模型来处理任意尺度的图像,但实际上它只用了 3 个尺度。只用 3 个尺度并不能说明模型能够处理任意多尺度的图像,这可能仅仅是暴力拟合的结果。因此,我觉得一个比较重要的思考方向是:怎么让去噪扩散模型更好地理解图像尺度这一概念,复用各个尺度的知识,从而实现任意多尺度的图像去噪。

受到经典图像表示方法拉普拉斯金字塔(Laplacian Pyramid)的启发,英伟达最近公布了一种叫做 Laplacian Diffusion Model (拉普拉斯扩散模型,后文简称 LaDM)的新型像素空间扩散模型,并用这种架构实现了文生图、超分辨率、ControlNet 等多种任务。在这篇博文里,我们来着重学习一下这种新型扩散模型的设计思想。

以往工作

扩散模型奠基之作 DDPM 及其升级版 ADM (Diffusion Models Beat GANs on Image Synthesis) 都是像素空间里的扩散模型。相比 LDM (隐扩散模型,即 Stable Diffusion),这类扩散模型不需要额外的自编码器来压缩图像,避免了编码解码带来的精度损失。

将图像从分辨率的维度拆解是一种很常见的思想。比如 Cascaded Diffusion Models 就是一种先生成低分辨率图像,再不断超分的扩散模型。今年比较有名的 VAR(Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction)也是一种按分辨率自回归的生成模型。

和这篇工作非常相关的早期工作是苹果在 2022 发表的 f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation。f-DM 将扩散模型的加噪推广到了降采样、模糊等其他退化策略上。降采样版的 f-DM 有非常多的设计和本工作很像。苹果该团队次年发表的 Matryoshka Diffusion Models 也用到了按分辨率逐次生成的设计。

将拉普拉斯金字塔融入扩散模型

拉普拉斯金字塔是一种图像表示方法,它把图像按频率成分拆成几张分辨率不同的图像,分辨率越低的图像表示频率越低的图像成分。我们直接通过下面的例子学习它的原理。假如 x 是原图,那么 x(3)=down(down(x))x(2)=down(x)-up(x(3)), x(1)=x-up(down(x))。对 x(1), x(2), x(3) 求加权和就可以还原输入图像。

受此启发,LaDM 将扩散模型的训练过程也用类似的方法分解:设 $\mathcal{X}$ 为训练图片集合,$\mathcal{X}^{(1)},\mathcal{X}^{(2)}, \mathcal{X}^{(3)}$ 分别是拉普拉斯金字塔不同成分构成的集合,那么我们在 $\mathcal{X}^{(3)}$, $\mathcal{X}^{(3)} \cup \mathcal{X}^{(2)}$,$\mathcal{X}$ 上分别训练三个去噪模型。也就是说,不同分辨率的模型生成不同层级的拉普拉斯金字塔复原结果。

根据经验,扩散模型早期(加噪后期)生成低频内容,后期(加噪前期)生成高频内容,所以训练时我们让不同分辨率的输入图像随噪声的衰退速度也不同。图像所代表的频率越低,衰减速度越慢,越需要从早期开始去噪。这样,在生成时,我们能生成到中途后再逐渐加上高频细节。

在采样过程中,我们按照下图所示的路线从低频到高频生成图像。有了该分辨率的初始图像后,按正常 DDPM 采样的步骤就可以生成当前分辨率的图像了。问题在于某分辨率的初始图像怎么从上一个分辨率过渡而来。

在切换当前带噪图像的分辨率时,我们既要放大其中的清晰图像(信号),也要放大其中的噪声。观察上一张图和下面的图,在分辨率切换时,新的高频成分(上图中的$\mathbf{x}^{(2)}$在时刻 3 及 $\mathbf{x}^{(2)}$ 在时刻 2)是一张纯黑图,新信号为零,所以对于信号的部分我们可以直接放大。而放大噪声时,我们要做一些噪声强度上的修改,保证放大后信噪比不变。这部分的细节详见论文。

1K 图像生成

为了生成 $1024 \times 1024$ 分辨率的图像,LaDM 采用了两阶段 Cascaded Diffusion Model 的设计,让生成高分辨率的图像约束于低分辨率图像。另外,由于注意力操作的时间复杂度很高,一般的像素扩散模型只能做到 $256 \times 256$ 大小。为了解决此问题,LaDM 依然用一个 $256 \times 256$ 的去噪模型来生成 1K 图片,但输入前后用小波变换来压缩/复原图像。

批判性分析与总结

这篇文章是一篇由公司发表的技术报告,展示了很多可视化结果,却没有任何定量结果,代码也没有开源,不知道它的生成能力和其他模型比起来如何。

这篇文章提出的模型虽然是像素空间扩散模型,但是其拉普拉斯金字塔的设计与模型是像素空间模型还是隐空间模型无关。我们完全可以把这套设计搬到隐空间上。VAR 已经向我们证明了对隐空间图像做拉普拉斯分解是可行的。另外,这篇文章的主干网络是 U-Net 而不是 DiT。想对这个工作做一点简单的改进的话,可以弄一个 LDM + DiT 版本的。

LaDM 设计最巧妙的点是其加噪过程,频率越高的成分越早变成纯噪声。这样的话我们可以在图像生成到一半的时候再直接把高频成分加上。如果高频成分一直在的话,我们还需要额外的设计在切换分辨率时把缺少的高频加上。

有工作证明神经网络不擅长拟合高频信息。因此,在图像任务中,手动将输入图像拆成不同频率成分可能有助于网络的学习。我们可能可以沿着这个思路去改进之前多种图像任务的输入。

随着视觉主干模型不断向 Transformer 靠拢,和 Transformer 配套的一些技术也从 NLP 社区涌入了 CV 社区。比如 Stable Diffusion 3 还在用标准 Transformer 那一套正弦位置编码,而其升级版 FLUX.1 就用上了旋转位置编码 (RoPE) , Lumina-T2X 模型甚至把 RoPE 的长度外推技术也从 NLP 社区搬了过来。在这篇博文中,我将站在一个对 NLP 技术了解不深的 CV 研究者的视角,较为详细地介绍一下 NLP 中 RoPE 相关的位置编码知识、RoPE 长度外推技术以及它们在 CV 里的应用。

长度外推,指的是使用在短序列上预训练的 Transformer 模型直接生成超出训练长度的长序列。类比到图像生成中,长度外推可以看成对模型所建模的图像分布做了一次超分辨率:比如模型训练时只见过 $256 \times 256$ 的图像,我们想直接用它生成 $512 \times 512$ 且同样清晰的图像。

推荐大家在阅读本文前先熟悉位置编码的基本原理,强烈推荐阅读 RoPE 提出者苏剑林的系列文章。

位置编码设计原则与 RoPE 的首次提出:https://kexue.fm/archives/8130

详细介绍 RoPE:https://kexue.fm/archives/8265

介绍长度外推的一项关键改进 (NTK-aware):https://kexue.fm/archives/9675 https://kexue.fm/archives/9706 https://kexue.fm/archives/9948

和这篇博文相关的两篇学术论文是:

YaRN,一种公认效果较好的长度外推技术:YaRN: Efficient Context Window Extension of Large Language Models (https://arxiv.org/abs/2309.00071)

Lumina-Next,前沿扩散 Transformer (Diffusion Transformer, DIT) 模型,采用了长度外推技术:Lumina-Next : Making Lumina-T2X Stronger and Faster with Next-DiT (https://arxiv.org/abs/2406.18583)

位置编码知识回顾

Transformer 中的位置编码

相比于此前流行的 CNN、RNN 模型,Transformer 的一大特点是其输出与输入次序无关。比如我们用 Transformer 建模文本的概率,那么模型会把「上海」和「海上」当成概率一样的词语。这也就是说 Transformer 无法从输入词元 (token) 的位置关系中获取信息。

如果让 Transformer 不输出信息聚合后的概率,还是保留输入词元的结构的话,那么打乱输入词元顺序就会同样地打乱输出词元顺序。模型依然无法获取输入词元间的位置关系。

为了把 1, 2, 3, 4 这样的位置信息输入进模型,标准 Tranformer 的做法是给不同位置的输入加上不同的位置编码。假设模型的中间变量都是二维向量,那么在句子中位置为 $k$ 的词元的位置编码是:

如果模型的中间变量都是 $d$ 维向量 (为了方便不妨认为 $d$ 是偶数),我们只需要把 $d$ 拆成 $d/2$ 组,每组用不同频率的三角函数即可。这样,长度为 $d$ 的词元在位置 $k$ 的位置编码是:

直观上来看,随着中间变量的维度越来越长,位置编码中对应的三角函数的频率不断变低,从一开始的 $1$ 逐渐靠向 $1/10000$。

上述公式来自论文,代码实现时要注意更多细节。比如有些代码中 $i$ 是从 $0$ 开始计数的。由于指数的分子的范围是从$0$到$d-2$,代码会把指数的分母也改成 $d-2$,保证最后一组三角函数的频率是 $1/10000$。

算出一个和输入词元向量等长的位置编码后,该编码会直接加到输入向量上。由于这种编码用了正弦函数,所以它被后续工作称为正弦位置编码。

相对位置编码与 RoPE

在设计位置编码时,最好能让编码传达词元的绝对位置相对位置信息。比如句号会出现在文本结尾而不是文本开头,这一规律来自绝对位置信息;而每几个词元会组成固定的词组,与它们在整段文本中的位置无关,这反映了相对位置信息的意义。

正弦位置编码同时满足了这两个性质。首先,正弦位置编码的输入只有绝对位置,它本质上就是一种绝对编码。另外,根据三角函数和角公式,假设偏移 $\Delta k$ 是常数,那么$sin(\alpha-\Delta k)$可以由含 $\alpha$ 的三角函数的线性组合表示。这说明模型能够从正弦编码中部分了解到一些相对位置的信息。

作为绝对位置编码,正弦编码虽然能够表达一些相对信息,但是这些信息太隐晦了。并且,该编码只在输入时加入,可能在网络运算中途这些信息就消失了。我们能不能更加显式地用某种绝对位置编码建模相对位置关系呢?

在 Transformer 中,不同位置的词元仅会在注意力操作时做信息交互。观察下面的注意力计算公式,更具体一点来说,信息交互发生在注意力的 QK 内积时。我们可以在每次注意力操作前都给 $Q$, $K$ 里各个向量加上位置编码,保证相对位置信息能反映在注意力计算里。

苏剑林设计了一种新的位置编码:考虑 $Q$ 里位置为 $m$ 的向量 $\mathbf{q}_m$ 和 $K$ 里位置为 $n$ 的向量 $\mathbf{k}_n$,给位置为 $j$ 的向量右乘上复数 $e^{ij\theta}$,其中 $i$ 是虚数单位, $\theta$ 是一个角度。这样,复数下的 QK 内积结果为:

其中,$Re[]$是取实部,$*$ 为共轭复数。可以发现,内积结果也出现了位置编码 $e^{i(m-n)\theta}$,且编码的值仅取决于相对位置 $(m-n)$。因此,这种编码能够更加显式地在注意力运算里建模相对位置。

由于最终结果取了实部,所以上述所有运算都可以转换成实数域的操作。假设 $\mathbf{q}_m = (q_0, q_1)$ 只是一个二维向量,那么上述右乘位置编码的操作可以写成:

从几何意义上讲,这个操作其实是二维向量旋转。因此,这种位置编码被称为旋转位置编码(RoPE)。

和正弦编码类似,要把二维中间变量拓展成 $d$ 维时,只要分组讨论,改变每组的频率(这里的频率是角度 $\theta$)就行了。

角度 $\theta$ 的设计可以参考 Transformer 的正弦编码:对于第 $i$ 组,我们令 $\theta_i=1/10000^{2i/d}$

RoPE 和正弦编码有同有异。相同之处在于:

  • 二者都是绝对位置编码,并通过编码公式的某些设计间接传递了相对位置信息。
  • 二者用了同样的正弦编码方式:随着变量通道数的增大,对应位置编码的正弦函数的频率不断指数衰减。

不同之处在于:

  • 正弦编码仅在模型输入时施加一次,RoPE 在所有自注意力计算时都施加。
  • 正弦编码会生成一组编码向量,加到输入上。而 RoPE 是一种操作,它的几何意义是把注意力输入向量旋转一个角度。

用 RoPE 实现长度外推

现在,我们正式进入本文的正题:长度外推。长度外推严格来说是一类任务,并不一定要用外推的做法。它似乎最早出自论文 ALiBi (Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation)。正如论文标题所示,该任务的目的就是「短训练,长推理」:在短序列上训练后,不经额外训练或只需少量微调,让模型生成长文本。后来这种任务也被称为「上下文窗口拓展」(Context Window Extension),目的依旧是用已经训好的模型来生成更大的文本,只是不强调方法是外推。为了称呼方便,我们在这篇博文里将该任务统称为「长度外推」。

我们想一想,假设模型训练时最大文本长度是 $L$,现在要生成长度为 $L’$ 的句子 ($L’>L$),我们需要做什么呢?其实我们只要把代码写好,除了生成长度以外啥也不改就行了。

这样的话,模型在运行时究竟哪里发生了变化呢?根据我们之前的分析,Transformer 是不知道位置信息的,只有位置编码传递了位置信息。因此,增加了生成句子长度后,原本只见过位置在 $L$ 之内的位置编码,现在要尝试解读位置为 $L, L+1, …, L’ - 1$ 的位置编码。因此,如果除了修改生成长度外什么也不做,其实就是让模型把学到的位置编码知识外推。但很可惜,由于没学过这些训练集之外的位置关系,这种外推法效果很差。

我们在接下来的几节里会讨论一些更加强大的长度外推策略。这里先补充介绍一点东西。看了对长度外推任务的基本介绍,读者或许会疑惑:长度外推似乎只要考虑位置编号就行了,不是非得和 RoPE 绑定起来?其实,长度外推真正要考虑的是位置编码的形式而不是只考虑编号。我们稍后的分析其实对所有类正弦编码都有效。但现在大家都是基于 RoPE 讨论,用基于 RoPE 的模型做实验,可能是因为 RoPE 更加直接、全面地建模了词元间的交互关系,只要调整了 RoPE 的公式,其效果立刻就能反映出来。相比之下,正弦位置编码只是在输入时提供了位置信息,修改位置编码的细节不能全面地影响模型的输出。

位置内插

既然超出 $L$ 的位置没有被训练过,那么在 $L$ 之内多选一些位置为分数的点不就行了?位置内插(Positional Interpolation, PI)就是这样的一种长度外推方法,它把长度为 $0$~$L’$ 的位置线性压缩到 $0$~$L$ 内。也就是说,对于位置$m$,将其的位置编号修改为:

由于位置编号会被送进正弦函数里,所以编号哪怕是分数也没关系。通过这种简单的线性内插方法,我们就能在已经学好的编号范围内多选一些位置,实现长度外推。

内插确实比外推的效果要好得多。后续所有长度外推方法实际上都是在研究如何更好的求插值位置编码。很快,有人就从频率分析的角度提出了线性内插的一个改进。

改变正弦函数基础频率:NTK-aware Scaled RoPE

直观认识

就在位置内插提出不久,就有研究者在社区 (https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/) 提出了一种效果更好,完全不需要微调的长度外推技术:NTK-aware Scaled RoPE (后文简称为”NTK-aware RoPE”)。该研究者后续将此方法进一步整理优化,发表了论文 YaRN: Efficient Context Window Extension of Large Language Models。我们先看一下 NTK-aware RoPE。

NTK-aware RoPE 的改动非常简洁,但它改动的地方却很出人意料:原来,总长度为 $d$ 的向量的第 $i$ 组位置编码的频率为:

现在,我们把 $10000$ 改掉,公式变为:

使用这种新长度外推方法,在上下文窗口大小为 2048 的 LLaMA 模型上,模型生成长文本的误差远低于之前的方法。

为什么这么简洁而奇怪的修改这么有效呢?在深入理解其原理之前,我们先直观地看一看这个方法具体修改了公式里的哪些参数。

先看新位置编码向量 $\sin m\theta_i^{NTK-aware}$ 的第二组(假设 $i$ 从 $0$ 开始计数),也就是含参频率最大的这一组。它现在是:

$\frac{L’}{L}$ 表示新长度是训练长度的几倍,它是一个大于 $1$ 的数。$d$ 在 LLaMA 里是 $128$,所以我们可以认为 $(d-2)$ 远大于 $2$。所以,$(\frac{L’}{L})^{(-2)/(d-2)}$ 这一项略小于 $1$。整体上看,这一改动差不多就是给三角函数的频率乘上了一个略小于 $1$ 的常数,几乎没变。

再看频率最小的 $i=d/2-1$ 项。它现在是:

而在线性内插中,我们直接把所有 $m$ 替换成了 $\frac{L}{L’}m$。所以,频率最小的项的公式和线性内插时的公式完全相同。

这里要澄清一下「外推」和「内插」的概念,这两个词的意义在很多博客和论文里并没有讲清楚。「内插」指的是通过像前面的线性位置内插一样,修改位置编号,使其恰好落在训练长度内。然而,一旦这个内插不够彻底,那么新位置编号就可能会超出训练长度,形成位置「外推」。我们本文讨论的所有方案,都是让不同频率的项在完全内插(恰好长度适合)和完全外推直接找一个平衡。一旦内插不彻底,就可以称为外推。所以很多文章里的「外推」,有的时候指的是不完全的内插。根据这样的术语定义,NTK-aware RoPE 的行为可以称为:最低频内插,其他频率外推。

从上面的分析可以看出,NTK-aware RoPE 还是沿用位置线性内插的思路,但是它对 RoPE 的影响更加平滑:对于位置编码高频项,公式几乎不变;对于最低频项,公式完全等于线性内插时的公式。

那么,NTK-aware RoPE 为什么有效呢?它又是怎么被想出来的呢?说起来,这个一直出现的 “NTK” 又是什么意思?NTK 其实是和神经网络相关的一种理论。NTK-aware RoPE 的提出者在构思这些公式时受到了 NTK 的启发,但他后续在论文里解释此方法时完全没有从严谨的理论入手,而只是讲了一些直觉的观察。在之后的两小节中,我将先从 NTK 理论的角度试图还原提出者的心路历程,再从一个广为人知、更易理解的角度来介绍 NTK-aware RoPE。

从 NTK 角度的解释

近几年和 NTK 理论比较相关的论文叫做 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains。这篇论文用 NTK 理论分析了 NeRF 这种以位置坐标为输入的 MLP 需要位置编码的原因,并将这类位置编码归纳为「傅里叶特征」。

这篇论文最大的一个发现是:在形式为 $[\sin \theta_i x, \cos \theta_i x]$ 这样的傅里叶特征中,最重要的是决定最大频率。最大频率越大,MLP 拟合高频信息的能力越强。

由于 RoPE 的公式来自于正弦位置编码,而正弦编码又可以看成一种特殊的傅里叶特征,所以 NTK-aware RoPE 的提出者也试图将傅里叶特征中的规律套用在 RoPE 上。他可能观察到了应用线性内插后 RoPE 公式(正弦编码公式)的频率变化。原来编码第 $i$ 项为:

应用线性内插后,公式为:

这里 $\frac{L}{L’}$ 是一个小于 $1$ 的数。所以,加上线性内插后,所有项的频率都变小了。自然,公式能表达的最大频率也变小了,拟合高频信息的能力下降了。

我们可以把线性内插类比到 NeRF 这类任务中。如果我们增加输入坐标的密度,确实可以让图片/3D 模型的输出分辨率变大。但是,根据信号处理的知识,这种分辨率变大并不能超出原有的频率,所以变大后的图片/3D 模型会看起来很模糊。「模糊」在文本任务中的体现可能就是误差指标上升。

出于这些原因,NTK-aware RoPE 的策略是尽可能不动高频项的频率,仅动低频项的频率。当然,按照这种设计思路,我们其实可以提出各种各样的方案。NTK-aware RoPE 选了实现起来最方便的一种:修改频率基底,让它在最低频时和线性内插对齐(读者感兴趣可以设方程自行推导频率基底的修改值,把我们刚刚有关最低频项的分析倒过来)。这样,自然就有高频项几乎不变,低频项向线性内插靠拢,也就是我们在上一小节中的观察。

根据我的理解,傅里叶特征本身就只是稍微用到 NTK 相关的理论(参见我有关傅里叶特征的博文)。而 NTK-aware RoPE 的作者貌似仅是受到了傅里叶特征的某些启发,完全没有严谨地用 NTK 理论来推导 NTK-aware RoPE 的形式。所以,我认为,要学习 NTK-aware RoPE,完全不用学习 NTK 理论。

NTK-aware RoPE 的提出者在互联网上和 YaRN 论文中用了一些更好理解的方式解释 NTK-aware RoPE。类似地,从进制转换的角度,苏剑林也发表了两篇一针见血的解读博文:https://kexue.fm/archives/9675 https://kexue.fm/archives/9706 。我建议从这些角度来学习 NTK-aware RoPE,然后忘掉 NTK 这个词。我们在下一节里就从这个角度重新认识一遍位置编码。

从进制的角度解释

其实几乎每个人都理解位置编码。

不信?我来问个问题:看到 $1234$ 后,你看到了几个数?

确实,这只是一个数。但是,我们人在看到这个数的时候,其实是看到了 $4$ 个十进制数字。通过把不同位置的数字组合,我们才理解了这个数究竟是多少。真正的数是一个概念,我们可以把两个东西这一概念,表示成汉字「二」,阿拉伯数字「2」,或者是二进制下的 $10$。我们常见的十进制只是表达数的一种方式。

而进制表示其实就是一种表达数的位置编码。想象一个十进制计时器,它的数字从 $0, 1, …$ 开始不断增长。每隔 1 次,个位变一次;每隔 10 次,十位变一次;每隔 100 次,百位变一次……。也就是说,个位是频率最高的,位数越高频率越低。是不是这和正弦位置编码很像?正弦位置编码和进制表示的区别在于,进制用求余体现周期性,正弦位置编码用正弦函数体现周期性。

长度外推,就好像一个只见过 0-999 的模型,突然要处理 1000 以上的数一样。为了只用三位数来表达更大的数,一种简单的做法是进制转换。比如我们直接把十进制变成十六进制,那么可以表达的数就从 $10^3$ 变成了 $16^3$。

回到正弦编码的公式里,进制这个概念体现在哪呢?进制的底数又是什么呢?

在十进制里,不同位表示十、百、千……每算一个更高的位的值,就要多除以一次 $10$。所以,在正弦编码里,我们需要关注哪个被除以的量在做指数运算。通过观察发现,正弦编码的底数是 $10000^{2/d}$。

知道了我们想把句子长度拓展几倍,我们就可以精确地算出新底数。通过这种方式,我们就能推导出 NTK-aware RoPE。也就是说,NTK-aware RoPE 修改频率基底其实就是对正弦函数做进制转换。这部分推荐大家去阅读前面提到的苏剑林的博文。

基于数字进制,我们可以把位置编码类比成表示时间的时钟,便于后续概念的理解。这是因为:

  • 正弦函数本身就可以用周期旋转来解释。
  • 相比数字的进制,时间的进制的底数是不同的:1 天有 24 个小时,而一小时有 60 分钟。这提示我们:我们不一定要对每种频率做同样的处理。

利用这个时钟的比喻,NTK-aware RoPE 的提出者在社区解释了不应该像线性内插一样修改最高频率的原因:就像我们用秒针来区分最精确的时间一样,神经网络用最高频的正弦编码区分相对位置关系,且只能看清 1 秒以上的偏差。使用线性内插后,最小的时间偏差是 0.5 秒,神经网络就不能很好地处理最高频的那块信息了。而 NTK-aware RoPE 不会修改一秒的定义,只会在分钟、小时等更低频的分量上多插值一点,神经网络依然能区分最精细的时间。

改进 NTK-aware RoPE:分部 NTK

我们在上一节中学到,NTK-aware RoPE 的设计思想是高频不动(或理解成高频外推),只对低频内插。只改频率基底虽然做法简洁,但不见得是最优的做法。高频不动这部分应该没什么问题,我们把目光放在 RoPE 的低频分量上。

还是从十进制的角度看待位置编码。假设训练集的位置只有 $0$~$2800$,那么在千位上,模型只见过 $0, 1, 2$ 三个数字。由于在千位上模型没有完整见过 $0$~$9$ 的循环,模型不能推测出其他几个数字的意义。因此,在千位上做长度外推时,一定要用内插把位置编号正确缩放到已学习的范围内。

这套分析怎么迁移到正弦编码上呢?对于十进制数字,我们能很快判断出某一位是否走完了一个周期。比如要把千位上的 $0$~$9$ 都走一遍,就至少得要一万个数。怎么找出正弦编码每个频率走一个周期需要的距离呢?

在正弦函数中,我们可以用 $2 \pi$ 除以频率,得到波长。正弦位置编码某一项的波长表示当训练上下文长度至少为多少时,这一项会「转」完一个周期。比如时钟上,秒针 60 秒转一圈,分针 3600 秒转一圈。

$2 \pi$ 除以频率明明算出的是周期,周期乘上速度才是波长。但 YaRN 的作者就是在论文里把这个量定义成了周长。可能他们认为波长的单位是长度,上下文窗口大小也是长度,两个单位是匹配的。我认为这个名字取得很糟糕,就应该叫做周期的,只不过周期的单位也是长度而已。

根据这个定义出来的波长,我们可以对正弦位置编码的不同位置分类讨论:

  • 如果波长过大,大于了训练时的文本长度,那么就用普通的线性内插,保证不在这些维度上外推。设它们的内插程度为 $1$。相比之下, NTK-aware RoPE 只对最低频项做了完整内插,而没有考虑其他波长过大的项也应该完整内插。
  • 如果波长过小,说明频率很高,不应该做任何修改。设它们的内插程度为 $0$。
  • 对于其他位置,根据它们的波长,线性选择内插程度。

这里波长过大、过小的阈值用超参数来决定,每个模型都需要手动调整。

总之,NTK-aware RoPE 只是模糊地定义了高频分量应该尽可能不变,低频分量应该尽可能像线性内插。而分部 NTK 则允许我们显式对各个频率分量做分类讨论。最终的位置编码方案 YaRN 在分部 NTK 的基础上还做了少许修改,对此感兴趣的读者可以去阅读论文。

图像生成中的 RoPE 与长度外推

了解了近年来 NLP 社区的位置编码技术,我们来以 Lumina-T2X 为例,再看一下这些技术是怎么用到视觉生成任务上的。

多维 RoPE

RoPE 本来是设计给 1D 的文本数据的。而在视觉任务中,图像是二维的,视频是三维的,我们需要设计更高维的位置编码。

回顾 RoPE 的形式:

要把它拓展成高维很简单。比如要拓展成 3D RoPE,只要把上面的公式复制两份,放到原公式的下面就行。也就是说,我们把向量拆成三份分别处理,每一部分和 1D RoPE 一样。

在这种设计下,模型所有中间向量的不同维度有了不同的意义,它们可能负责了视频宽度、高度或长度上的信息处理。我们也可以根据实际需要,让负责不同视频维度的向量长度不同。

视觉扩散模型中 RoPE 的长度外推设计

为了生成比训练分辨率更大的图像,Lumina-T2X 也参考了 NTK-aware RoPE,提出了一些和图像相关的 RoPE 改进策略。

首先,和分布 NTK 策略一样,Lumina-T2X 提出了频率感知 RoPE。在这种策略下,波长大于等于训练长度的位置编码项完全使用线性内推,剩下的项使用 NTK-aware RoPE。

另外,Lumina-T2X 还提出了时刻感知 RoPE。这个「时刻」指的是扩散模型里的加噪/去噪时刻。根据实验结果,Lumina-T2X 的作者发现线性内插会保持图像整体结构,但是图像局部质量下降;NTK-aware 策略提升了局部质量,却会出现内容重复现象,也就是全局关系不合理。能不能在某一方面结合二者呢?根据之前使用扩散模型的经验,扩散模型在去噪初期只生成低频信息,也就是全局信息,后期才会生成高频细节。受此启发,Lumina-T2X 提出了时刻感知 RoPE,该策略会在去噪早期仅使用线性内插,后续慢慢过渡到频率感知 RoPE。

以下是论文展示的在各种长度外推策略下生成 2K 图片的效果图。最左侧的 1K 图片供参考。

总结

长度外推是生成任务中的一项重要技术,它让我们在不大规模重新训练模型的前提下提升输出内容的长度/大小。而 Transformer 本身是一种无法获取输入元素位置信息的生成模型,需要靠额外的位置编码来感知位置。那么正好,只要我们能够适当地修改位置编码的推理行为,就能想办法让模型生成更长的内容。目前长度外推的方案都和修改 RoPE——一种给 Transformer 显式提供相对位置信息的位置编码——有关。我们主要学习了 NTK-aware RoPE 的设计原理,并通过深入的分析学习了其改进版分部 NTK RoPE。基于这些知识,我们简单认识了 RoPE 长度外推在视觉生成中的应用,其中比较有趣的一项设计是做长度外推时考虑扩散模型的去噪时刻。

说白了,本文所有长度外推设计都是在从两个维度上排列组合:RoPE 可以看成是由多个频率项组成的正弦编码;外推方案可以从位置编号线性内插过渡到位置编号不变(即位置外推)。一般的设计策略是:对于没有学满一个完整周期的频率项,采用完全线性内插;对于其余频率项,按一定比例执行线性内插。加上了扩散模型的去噪时刻这一设计维度后,我们可以按同样的设计思路,早期更关注低频,晚期更关注高频。

我觉得长度外推技术的能力是有上限的。我们完全可以从信号处理或者信息论的角度来思考这一问题,因为它的本质和从频域对图像做超分辨率很像。在较短的序列中,模型只能学到这种长度的序列所能表示的最大频率的信息。强行用它来生成更长的序列,只会出现两种情况:要么序列局部不够清晰,要么每个局部很清晰但是没有很好的全局依赖关系。根据信息论,模型就是不能从短序列中学到长序列蕴含的一些规律。从 Lumina-T2X 展示的结果里,我感觉 NTK-aware RoPE 的做法某种程度上就像是把全图做超分辨率变成拆成几个小图,每个小图在原来的训练长度上分别做超分辨率。这样最后图像每一块都很清晰,但合起来看就有问题。可能对于一些文本任务来说,只要局部质量高就行了,长距离依赖没那么重要。

最近我在看位置编码最新技术时,看到了一个叫做 “NTK-aware” 的词。我想:「”NTK”是什么?Next ToKen (下一个词元)吗?为什么要用这么时髦的缩写?」看着看着,我才发现不对劲。原来,NTK 是神经网络理论里的一个概念,它从 kernel regression 的角度解释了神经网络的学习方法。基于 NTK 理论,有人解释了位置编码的理论原理并将其归纳为一种特殊的 Fourier Feature (傅里叶特征)。这么多专有名词一下就把我绕晕了,我花了几天才把它们之间的关系搞懂。

在这篇文章里,我主要基于论文 Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains (后文简称为「傅里叶特征论文」),介绍傅里叶特征这一概念。为了讲清这些理论的发展脉络,我会稍微讲一下 NTK 等理论概念。介绍完傅里叶特征后,我还会讲解它在其他方法中的应用。希望读完本文后,读者能够以这篇论文为基点,建立一个有关位置编码原理的知识网络,以从更深的层次来思考新的科研方向。

用 MLP 表示连续数据

我们先从一个具体的任务入手,直观体会傅里叶特征能够做些什么事。

我们知道,神经网络,哪怕是最简单的多层感知机(MLP),都有着很强的泛化能力:训练完毕后,对于训练集里完全没见过的输入,网络也能给出很正确的输出。特别地,如果新输入恰好和训练集的某个输入很近,那么它的输出也会和对应的训练集输出很近;随着新输出与训练集输入的距离不断增加,新输出也会逐渐变得不同。这反映了神经网络的连续性:如果输入的变化是连续的,那么输出的变化也是连续的。

基于神经网络的这一特性,有人想到:我们能不能用神经网络来表示连续数据呢?比如我想表达一张处处连续的图像,于是我令神经网络的输入是 (x, y) 表示的二维坐标,输出是 RGB 颜色。之后,我在单张图像上过拟合这个 MLP。这样,学会表示这张图像后,哪怕输入坐标是分数而不是整数,神经网络也能给出一个颜色输出。

这种连续数据有什么好处呢?我们知道,计算机都是以离散的形式来存储数据的。比如,我们会把图像拆成一个个像素,每个像素存在一块内存里。对于图像这种二维数据,计算机的存储空间还勉强够用。而如果想用密集的离散数据表达更复杂的数据,比如 3D 物体,计算机的容量就捉襟见肘了。但如果用一个 MLP 来表达 3D 物体的话,我们只需要存储 MLP 的参数,就能获取 3D 物体在任何位置的信息了。

这就是经典工作神经辐射场 (Neural Radiance Field, NeRF) 的设计初衷。NeRF 用一个 MLP 拟合 3D 物体的属性,其输入输出如下图所示。我们可以用 MLP 学习每个 3D 坐标的每个 2D 视角处的属性(这篇文章用的属性是颜色和密度)。根据这些信息,利用某些渲染算法,我们就能重建完整的 3D 物体。

上述过程看起来好像很简单直接。但在 NeRF 中,有一个重要的实现细节:必须给输入加上位置编码,MLP 才能很好地过拟合连续数据。这是为什么呢?让我们先用实验复现一下这个现象。

MLP 拟合连续图像实验

为了快速复现和位置编码相关的问题,我们简单地用一个 MLP 来表示图像:MLP 的输入是 2D 坐标,输出是此处的三通道 RGB 颜色。我为这篇博文创建一个 GitHub 文件夹 https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/FourierFeature ,该实验的 Notebook 代码在文件夹的 image_mlp.ipynb 中,欢迎大家 clone 项目并动手尝试。

一开始,我们先导入库并可视化要拟合的图片。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.io import read_image, ImageReadMode
from torchvision.transforms.functional import to_pil_image

from tqdm import tqdm
from einops import rearrange

def viz_image(pt_img: torch.Tensor):
pil_img = to_pil_image(pt_img)
display(pil_img)


input_image = read_image('misuzu.png', ImageReadMode.RGB)
input_image = input_image.to(torch.float32) / 255
input_image = input_image.unsqueeze(0)
input_image = F.interpolate(input_image, (256, 256), mode='bilinear')
viz_image(input_image[0])

我们再定义一个 MLP 类。稍后我们会并行地传入二维坐标。具体来说,我们会将输入定义为一个 [1, 2, H, W] 形状的数据,其中通道数 2 表示 (i, j) 格式的坐标。由于输入是以图像的形式并行输入的,我们可以用 $1 \times 1$ 的 2D 卷积来表示二维数据上的并行 MLP。所以在下面这个 MLP 里,我们只用到 $1 \times 1$ 卷积、激活函数、归一化三种层。按照傅里叶特征论文的官方示例,网络最后要用一个 Sigmoid 激活函数调整输出的范围。

1
2
3
4
5
6
7
8
9
10
11
12
class MLP(nn.Module):
def __init__(self, in_c, out_c=3, hiden_states=256):
super().__init__()
self.mlp = nn.Sequential(
nn.Conv2d(in_c, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),
nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),
nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),
nn.Conv2d(hiden_states, out_c, 1), nn.Sigmoid()
)

def forward(self, x):
return self.mlp(x)

之后我们来定义训练数据。在一般的任务中,输入输出都是从训练集获取的。而在这个任务中,输入是二维坐标,输出是图像的颜色值。输出图像 input_image 我们刚刚已经读取完毕了,现在只需要构建输入坐标即可。我们可以用下面的代码构建一个 [1, 2, H, W] 形状的二维网格,grid[0, :, i, j] 处的数据是其坐标 (i, j) 本身。当然,由于神经网络的输入一般要做归一化,所以我们会把原本 0~H0~W 里的高宽坐标缩放都到 0~1。最终 grid[0, :, i, j]==(i/H, j/W)

1
2
3
4
5
H, W = input_image.shape[2:]

h_coord = torch.linspace(0, 1, H)
w_coord = torch.linspace(0, 1, W)
grid = torch.stack(torch.meshgrid([h_coord, w_coord]), -1).permute(2, 0, 1).unsqueeze(0)

准备好一切后,我们就可以开始训练了。我们初始化模型 model 和优化器 optimizer,和往常一样训练这个 MLP。如前所述,这个任务的输入输出非常直接,输入就是坐标网格 grid,目标输出就是图片 input_image。每训练一段时间,我们就把当前 MLP 拟合出的图片和误差打印出来。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = MLP(2).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
input_image = input_image.to(device)
grid = grid.to(device)
for epoch in tqdm(range(n_loops)):
output = model(grid)
loss = F.l1_loss(output, input_image)
optimizer.zero_grad()
loss.backward()
optimizer.step()

if epoch % 100 == 0 or epoch == n_loops - 1:
viz_image(output[0])
print(loss.item())

运行代码,大致能得到如下输出。可以看到,从一开始,图像就非常模糊。

不过,如果我们在把坐标输入进网络前先将其转换成位置编码——一种特殊的傅里叶特征,那么 MLP 就能清晰地拟合出原图片。这里我们暂时不去关注这段代码的实现细节。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class FourierFeature(nn.Module):
def __init__(self, in_c, out_c, scale):
super().__init__()
fourier_basis = torch.randn(in_c, out_c // 2) * scale
self.register_buffer('_fourier_basis', fourier_basis)

def forward(self, x):
N, C, H, W = x.shape
x = rearrange(x, 'n c h w -> (n h w) c')
x = x @ self._fourier_basis
x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)

x = 2 * torch.pi * x
x = torch.cat([torch.sin(x), torch.cos(x)], dim=1)
return x

feature_length = 256
model = MLP(feature_length).to(device)
fourier_feature = FourierFeature(2, feature_length, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
for epoch in tqdm(range(n_loops)):
x = fourier_feature(grid)
output = model(x)
loss = F.l1_loss(output, input_image)
optimizer.zero_grad()
loss.backward()
optimizer.step()

if epoch % 100 == 0 or epoch == n_loops - 1:
viz_image(output[0])
print(loss.item())
prev_output = output

简单地对比一下,此前方法的主要问题是 MLP 无法拟合高频的信息(如图块边缘),只能生成模糊的图像。而使用位置编码后,MLP 从一开始就能较好地表示高频信息。可见,问题的关键在于如何让 MLP 更好地拟合数据的高频信息。

接下来,我们来从一个比较偏理论的角度看一看论文是怎么分析位置编码在拟合高频信息中的作用的。

核回归

傅里叶特征论文使用了神经正切核(Nerual Tangent Kernel, NTK)来分析 MLP 的学习规律,而 NTK 又是一种特殊的核回归 (Kernel Regression) 方法。在这一节里,我会通过代码来较为仔细地介绍核回归。下一节我会简单介绍 NTK。

和神经网络类似,核回归也是一种数学模型。给定训练集里的输入和输出,我们建立这样一个模型,用来拟合训练集表示的未知函数。相比之下,核回归的形式更加简单,我们有更多的数学工具来分析其性质。

核回归的设计思想来源于我们对于待拟合函数性质的观察:正如我们在前文的分析一样,要用模型拟合一个函数时,该模型在训练数据附近最好是连续变化的。离训练集输入越近,输出就要和其对应输出越近。基于这种想法,核回归直接利用和所有数据的相似度来建立模型:假设训练数据为 $(x_i, y_i), i \in [1, n]$,我们定义了一个计算两个输入相似度指标 $K(x_1, x_2)$,那么任意输入 $x$ 的输出为:

也就是说,对于一个新输入 $x$,我们算它和所有输入 $x_i$ 的相似度 $w_i$,并把相似度归一化。最后的输出 $f(x)$ 是现有 $y_i$ 的相似度加权和。

这样看来,只要有了相似度指标,最终模型的形式也就决定下来了。我们把这个相似度指标称为「核」。至于为什么要把它叫做核,是因为这个相似度指标必须满足一些性质,比如非负、对称。但我们这里不用管那么多,只需要知道核是一种衡量距离的指标,决定了核就决定了核回归的形式。

我们来通过一个简单的一维函数拟合实验来进一步熟悉核回归。该实验代码在项目文件夹下的 kernel_regression.ipynb 中。

先导入库。

1
2
3
4
%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

再创建一个简单的非线性函数,做为我们的拟合目标。这个函数就是一个简单的周期为 $2$ 的正弦函数乘上线性函数 $(1-x)$。我们可以简单可视化一下函数在 $[-1, 1]$ 之间的图像。

1
2
3
4
5
6
7
def func(x):
return np.sin(np.pi * x) * (1 - x)

xs = np.linspace(-1, 1, 100)
ys = func(xs)
plt.plot(xs, ys)
plt.show()

基于这个函数,我们等间距地选一些点做为训练数据。

1
2
3
4
sample_x = np.linspace(-1, 1, 10)
sample_y = func(sample_x)
plt.scatter(sample_x, sample_y)
plt.show()

有了数据后,我们来用核回归根据数据拟合这个函数。在决定核回归时,最重要的是决定核的形式。这里我们用正态分布的概率密度函数来表示核,该核唯一的超参数是标准差,需要我们根据拟合结果手动调整。标准差为 1 的标准正态分布核的图像如下所示。由于最后要做归一化,正态分布密度函数的系数被省略掉了。

1
2
3
4
5
6
7
def kernel_func(x_ref, x_input, sigma=1):
return np.exp(-(x_input-x_ref)**2 / (2 * sigma**2))

xs = np.linspace(-1, 1, 100)
ys = kernel_func(0, xs)
plt.plot(xs, ys)
plt.show()

可以从图像中看出,离某输入越近(假设该输入是 0),那么相似度就越高。这符合我们对于相似度函数的要求。

有了核函数后,我们就直接得到了模型。根据核回归模型计算结果的函数为 kernel_regression。函数参数 xs, ys 表示训练数据,x_input 表示测试时用的输入坐标,sigma 是核回归的超参数。

假设有 n 个训练样本,有 m 个测试输入,那么我们要计算每个测试输入对每个训练输入的 n * m 个相似度,这些相似度会存到矩阵 weight 里。为此,我们需要对 xsx_input 做一些形状变换,再用上面定义的核函数 kernel_func 求出每对相似度。有了相似度后,我们根据公式计算点乘结果 weight_dot 及归一化系数 weight_sum,并最终计算出核回归的结果 res

基于这个函数,我们可以将测试输入定义成 [-1, 1] 上一些更密集的坐标,并用上面定义好的 10 个样本做为训练集,得到核回归的结果。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def kernel_regression(xs, ys, x_input, sigma=1):
# xs: [n, ]
# ys: [n, ]
# x_input: [m, ]
N = xs.shape[0]
xs = np.expand_dims(xs, 1)
ys = np.expand_dims(ys, 1)
x_input = np.expand_dims(x_input, 0)
x_input = np.repeat(x_input, N, 0)
weight = kernel_func(xs, x_input, sigma) # [n, m]
weight_sum = np.sum(weight, 0)
weight_dot = weight.T @ ys
weight_dot = np.squeeze(weight_dot, 1)
res = weight_dot / weight_sum
return res

sigma = 1
xs = np.linspace(-1, 1, 100)
ys = kernel_regression(sample_x, sample_y, xs, sigma)
plt.title(f'sigma = {sigma}')
plt.plot(xs, ys)
plt.show()

我们可以通过修改 sigma 来得到不同的拟合效果。以下是我的一些结果:

可以看出,标准差越小,模型倾向于过拟合;随着标准差变大,曲线会逐渐平缓。我们需要不断调整超参数,在过拟合和欠拟合之间找到一个平衡。这种现象很容易解释:正态分布核函数的标准差越小,意味着每个训练数据的影响范围较小,那么测试样本更容易受到少数样本的影响;标准差增大之后,各个训练样本的影响开始共同起作用,我们拟合出的函数也越来越靠近正确的函数;但如果标准差过大,每个训练样本的影响都差不多,那么模型就什么都拟合不了了。

从实验结果中,我们能大致感受到核回归和低通滤波很像,都是将已知数据的平均效果施加在未知数据上。因此,在分析核回归的时候,往往会从频域分析核函数。如果核函数所代表低通滤波器的带宽 (bandwidth)越大,那么剩下的高频信息就更多,核回归也更容易拟合高频信息较多的数据。

神经正切核

那么,核回归是怎么和神经网络关联起来的呢?有研究表明,在一些特殊条件下,MLP 的最终优化结果可以用一个简单的核回归来表示。这不仅意味着我们可以神奇地提前预测梯度下降的结果,还可以根据核回归的性质来分析神经网络的部分原理。这种能表示神经网络学习结果的核被称为神经正切核(NTK)。

这些特殊条件包括 MLP 无限宽、SGD 学习率的学习率趋近 0 等。由于这些条件和实际神经网络的配置相差较远,我们难以直接用核回归预测复杂神经网络的结果。不过,我们依然可以基于这些理论来分析和神经网络相关的问题。傅里叶特征的分析就是建立在 NTK 上的。

NTK 的形式为

其中,$f$ 是参数为 $\theta$ 的神经网络,$\langle\cdot,\cdot \rangle$为内积运算。简单来看,这个式子是说神经网络的核回归中,任意两个向量间的相似度等于网络对参数的偏导的内积的期望。基于 NTK,我们可以分析出很多神经网络的性质,比如出乎意料地,神经网络的结果和随机初始化的参数无关,仅和网络结构和训练数据有关。

在学习傅里叶特征时,我们不需要仔细研究这些这些理论,而只需要知道一个结论:一般上述 NTK 可以写成标量函数 $h_{NTK}(\mathbf{x}_i^T\mathbf{x}_j)$,也就是可以先算内积再求偏导。这意味用核回归表示神经网络时,真正要关心的是输入间的内积。别看 NTK 看起来那么复杂,傅里叶特征论文其实主要就用到了这一个性质。

为了从理论上讲清为什么 MLP 难以拟合高频,作者还提及了很多有关 NTK 的分析,包括一种叫做谱偏差(spectral bias)的现象:神经网络更容易学习到数据中的低频特征。可能作者默认读者已经熟悉了相关的理论背景,这部分论述经常会出现逻辑跳跃,很难读懂。当然,不懂这些理论不影响理解傅里叶特征。我建议不要去仔细阅读这篇文章有关谱偏差的那一部分。

正如我们在前文的核回归实验里观察到的,核回归模型能否学到高频取决于核函数的频域特征。因此,这部分分析和 NTK 的频域有关。对这部分内容感兴趣的话可以去阅读之前有关谱偏差的论文。

傅里叶特征的平移不变性

在上两节中,我们花了不少功夫去认识谱回归和 NTK。总结下来,其实我们只需要搞懂两件事:

  • 神经网络最终的收敛效果可以由简单的核回归决定。而核回归重点是定义两个输入之间的相似度指标(核函数)。
  • 表示神经网络的核回归相似度指标是 NTK,它其实又只取决于两个输入的内积 $\mathbf{x}_i^T\mathbf{x}_j$。

根据这一性质,我们可以部分解释为什么在文章开头那个 MLP 拟合连续图像的实验中,位置编码可以提升 MLP 拟合高频信息的能力了。这和位置输入的特性有关。

当 MLP 的输入表示位置时,我们希望模型对输入位置具有平移不变性。比如我们现在有一条三个样本组成的句子 $(1, A), (2, B), (3, C)$。当我们同时改变句子的位置信息时,比如将句子的位置改成 $(11, A), (12, B), (13, C)$时,网络能学出完全一样的东西。但显然不对输入位置做任何处理的话, $(1, 2, 3)$ 和 $(11, 12, 13)$ 对神经网络来说是完全不同的意思。

而使用位置编码的话,情况就完全不同了。假如输入数据是二维坐标 $\mathbf{v}\in [0, 1)^d$,我们可以用下面的式子建立一个维度为 $2m$ 的位置编码:

其中 $a_i$ 是系数, $b \in \mathbb{R}^{m \times 2}$ 是一个投影矩阵,用于把原来 2D 的位置变成一个更长的位置编码。当然,由于位置编码中既要有 $\sin$ 也要有 $\cos$,所以最终的位置编码长度为 $2m$。

根据我们之前的分析,NTK 只取决于输入间的内积。算上位置编码后,一对输入位置 $\mathbf{v}_1, \mathbf{v}_2$ 的内积为:

而根据三角函数和角公式可知:

这样,上面那个内积恰好可以写成:

上式完全由位置间的相对距离 $\mathbf{v_1}-\mathbf{v_2}$ 决定。上式决定了 NTK,NTK 又决定了神经网络的学习结果。所以,神经网络的收敛结果其实完全取决于输入间的相对距离,而不取决于它们的绝对距离。也因此,位置编码使得 MLP 对于输入位置有了平移不变性。

加入位置编码后,虽然 MLP 满足了平移不变性,但这并不代表 MLP 学习高频信息的能力就变强了。平移不变性能给我们带来什么好处呢?作者指出,当满足了平移不变性后,我们就能手动调整 NTK 的带宽了。回想一下我们上面做的核回归实验,如果我们能够调整核的带宽,就能决定函数是更加高频(尖锐)还是更加低频(平滑)。这里也是同理,如果我们能够调大 NTK 的带宽,让它保留更多高频信息,那么 MLP 也就能学到更多的高频信息。

作者在此处用信号处理的知识来分析平移不变性的好处,比如讲了新的 NTK 就像一个重建卷积核 (reconstruction filter),整个 MLP 就像是在做卷积。还是由于作者省略了很多推导细节,这部分逻辑很难读懂。我建议大家直接记住推理的结论:平移不变性使得我们能够调整 NTK 的带宽,从而调整 MLP 学习高频的能力。

那我们该怎么调整 NTK 的带宽呢?现在的新 NTK 由下面的式子决定:

为了方便分析,我们假设$\mathbf{v}$和$\mathbf{b_j}$都是一维实数。那么,如果我们令$b_j=j$的话:

这个式子能令你想到什么?没错,就是傅里叶变换。$j$ 较大的项就表示 NTK 的高频分量。我们可以通过修改前面的系数 $a_j$ 来手动调整 NTK 的频域特征。我们能看到,位置编码其实就是在模拟傅里叶变换,所以作者把位置编码总结为傅里叶特征。

作者通过实验证明我们可以手动修改 NTK 的频谱。实验中,作者令 $b_j=j, a_j=1/j^p$。$p=\infty$ 表示位置编码只有第一项:$\gamma(v)=[\cos 2\pi v, \sin 2\pi v]^T$。不同 $p$ 时 NTK 的空域和频域示意图如下所示。可以看出,令 $p=0$ 时,即傅里叶特征所有项的系数都为 $1$ 时,NTK 的高频分量不会衰减。这也意味着 MLP 学高频信息和低频信息的能力差不多。

随机傅里叶特征

现在我们已经知道傅里叶特征的公式是什么,并知道如何设置其中的参数 $a_j$, $\mathbf{b}_j$ 了。现在,还有一件事我们没有决定:该如何设置傅里叶特征的长度 $m$ 呢?

既然我们说傅里叶特征就是把输入的位置做了一次傅里叶变换,那么一般来讲,傅里叶特征的长度应该和原图像的像素数一样。比如我们要表示一个 $256 \times 256$ 的图像,那么我们就需要令 $m = 256 \times 256 / 2$ ,$\mathbf{b}$ 表示不同方向上的频率:$[(1, 1), (1, 2), …, (128, 256)]$。但这样的话,神经网络的参数就太多了。可不可以令 $m$ 更小一点呢?

根据之前的研究 Random features for large-scale kernel machines 表明,我们不需要密集地采样傅里叶特征,只需要稀疏地采样就行了。具体来说,我们可以从某个分布随机采样 $m$ 个频率 $\mathbf{b_j}$ 来,这样的学习结果和密集采样差不多。当然,根据前面的分析,我们还是令所有系数 $a_j=1$。在实验中,作者发现,$\mathbf{b_j}$ 从哪种分布里采样都无所谓,关键是 $\mathbf{b_j}$ 的采样分布的标准差,因为这个标准差决定了傅里叶特征的带宽,也决定了网络拟合高频信息的能力。实验的结果如下:

我们可以不管图片里 $1/f^x$ 是啥意思,只需要知道 a, b, c 是三组不同的实验就行。虚线是密集采样傅里叶特征的误差,它的结果反映了一个「较好」的误差值。令人惊讶的是,不管从哪种分布里采样 $\mathbf{b_j}$,最后学出来的网络误差都差不多。问题的关键在于采样分布的标准差。把标准差调得够好的话,模型的误差甚至低于密集采样的误差。

也就是说,虽然我们花半天分析了位置编码和傅里叶变换的关系,但我们没必要照着傅里叶变换那样密集地采样频率,只需要随机选一些频率即可。当然,这个结论只对 MLP 拟合连续数据的任务有效,和 Transformer 里的位置编码无关。

代码实现随机傅里叶特征

现在,我们可以回到博文开头的代码,看一下随机傅里叶特征是怎么实现的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class FourierFeature(nn.Module):
def __init__(self, in_c, out_c, scale):
super().__init__()
fourier_basis = torch.randn(in_c, out_c // 2) * scale
self.register_buffer('_fourier_basis', fourier_basis)

def forward(self, x):
N, C, H, W = x.shape
x = rearrange(x, 'n c h w -> (n h w) c')
x = x @ self._fourier_basis
x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)

x = 2 * torch.pi * x
x = torch.cat([torch.sin(x), torch.cos(x)], dim=1)
return x

feature_length = 256
model = MLP(feature_length).to(device)
fourier_feature = FourierFeature(2, feature_length, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
n_loops = 400
for epoch in tqdm(range(n_loops)):
x = fourier_feature(grid)
output = model(x)
loss = F.l1_loss(output, input_image)
optimizer.zero_grad()
loss.backward()
optimizer.step()

if epoch % 100 == 0 or epoch == n_loops - 1:
viz_image(output[0])
print(loss.item())
prev_output = output

傅里叶特征通过类 FourierFeature 实现。其代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class FourierFeature(nn.Module):
def __init__(self, in_c, out_c, scale):
super().__init__()
fourier_basis = torch.randn(in_c, out_c // 2) * scale
self.register_buffer('_fourier_basis', fourier_basis)

def forward(self, x):
N, C, H, W = x.shape
x = rearrange(x, 'n c h w -> (n h w) c')
x = x @ self._fourier_basis
x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)

x = 2 * torch.pi * x
x = torch.cat([torch.sin(x), torch.cos(x)], dim=1)
return x

构造函数里的 fourier_basis 表示随机傅里叶特征的频率,对应论文公式里的$\mathbf{b}$,scale 表示采样的标准差。初始化好了随机频率后,对于输入位置 x,只要按照公式将其投影到长度为 out_c / 2 的向量上,再对向量的每一个分量求 sin, cos 即可。按照之前的分析,我们令所有系数 $a$ 为 $1$,所以不需要对输出向量乘系数。

傅里叶特征在 StyleGAN3 里的应用

傅里叶特征最经典的应用就是 NeRF 这类过拟合连续数据任务。除此之外,傅里叶特征另一次大展身手是在 StyleGAN3 中。

StyleGAN3 希望通过平滑地移动生成网络的输入来使输出图片也发生对应的移动。为此,StyleGAN3 将生成网络的输入定义为频域上的一个有限带宽图像信号:根据信号处理知识,我们能够将有限带宽信号转换成空域上无限连续的信号。也就是说,不管输入的分辨率(采样率)多低,我们都能够平滑地移动输入图片。StyleGAN3 借助随机傅里叶特征来实现这样一个频域图像。

以下代码选自 StyleGAN3 中傅里叶特征的构造函数。这个函数的关键是随机生成一些频率固定,但方向可以不同的傅里叶频率。函数先随机采样了一些频率,再将它们归一化,最后乘上指定的带宽 bandwidth,保证所有频率大小相等。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class SynthesisInput(torch.nn.Module):
def __init__(self,
w_dim, # Intermediate latent (W) dimensionality.
channels, # Number of output channels.
size, # Output spatial size: int or [width, height].
sampling_rate, # Output sampling rate.
bandwidth, # Output bandwidth.
):
super().__init__()
self.w_dim = w_dim
self.channels = channels
self.size = np.broadcast_to(np.asarray(size), [2])
self.sampling_rate = sampling_rate
self.bandwidth = bandwidth

# Draw random frequencies from uniform 2D disc.
freqs = torch.randn([self.channels, 2])
radii = freqs.square().sum(dim=1, keepdim=True).sqrt()
freqs /= radii * radii.square().exp().pow(0.25)
freqs *= bandwidth
phases = torch.rand([self.channels]) - 0.5

而在使用这个类获取网络输入时,和刚刚的 MLP 实现一样,我们会先生成一个二维坐标表格 grid 用于查询连续图片每一处的颜色值,再将其投影到各个频率上,并计算新向量的正弦函数。

这段代码中,有两块和我们自己的实现不太一样。第一,StyleGAN3 允许对输入坐标做仿射变换(平移和旋转)。仿射变换对坐标的影响最终会转化成对三角函数相位 phases 和频率 freqs 的影响。第二,在计算三角函数时,StyleGAN3 只用了正弦函数,没有用余弦函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def forward(self, ...):
...

# Transform frequencies.
phases = ...
freqs = ...

# Construct sampling grid.
theta = torch.eye(2, 3, device=w.device)
theta[0, 0] = 0.5 * self.size[0] / self.sampling_rate
theta[1, 1] = 0.5 * self.size[1] / self.sampling_rate
grids = torch.nn.functional.affine_grid(theta.unsqueeze(0), [1, 1, self.size[1], self.size[0]], align_corners=False)

# Compute Fourier features.
x = (grids.unsqueeze(3) @ freqs.permute(0, 2, 1).unsqueeze(1).unsqueeze(2)).squeeze(3) # [batch, height, width, channel]
x = x + phases.unsqueeze(1).unsqueeze(2)
x = torch.sin(x * (np.pi * 2))
x = x * amplitudes.unsqueeze(1).unsqueeze(2)

...

# Ensure correct shape.
x = x.permute(0, 3, 1, 2) # [batch, channel, height, width]
return x

我们在 MLP 拟合连续图像的实验里复现一下这两个改动。首先是二维仿射变换。给定旋转角 theta 和两个方向的平移 tx, ty,我们能够构造出一个 $3 \times 3$ 的仿射变换矩阵。把它乘上坐标 [x, y, 1] 后,就能得到仿射变换的输出。我们对输入坐标 grid 做仿射变换后得到 grid_ext,再用 grid_ext 跑一遍傅里叶特征和 MLP。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
N, C, H, W = grid.shape
tx = 50 / H
ty = 0
theta = torch.tensor(torch.pi * 1 / 8)
affine_matrix = torch.tensor([
[torch.cos(theta), -torch.sin(theta), tx],
[torch.sin(theta), torch.cos(theta), ty],
[0, 0, 1]
]
).to(device)
grid_ext = torch.ones(N, 3, H, W).to(device)
grid_ext[:, :2] = grid.clone()
grid_ext = grid_ext.permute(0, 2, 3, 1)
grid_ext = (grid_ext @ affine_matrix.T)
grid_ext = grid_ext.permute(0, 3, 1, 2)[:, :2]

x = fourier_feature(grid_ext)
output = model(x)
viz_image(output[0])

在示例代码中,我们可以得到旋转 45 度并向下平移 50 个像素的图片。可以看到,变换成功了。这体现了连续数据的好处:我们可以在任意位置对数据采样。当然,由于这种连续数据是通过过拟合实现的,在训练集没有覆盖的坐标处无法得到有意义的颜色值。

之后,我们来尝试在傅里叶特征中只用正弦函数。我们将投影矩阵的输出通道数从 out_c / 2 变成 out_c,再在 forward 里只用 sin 而不是同时用 sin, cos。经实验,这样改了后完全不影响重建质量,甚至由于通道数更多了,重建效果更好了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class FourierFeature(nn.Module):
def __init__(self, in_c, out_c, scale):
super().__init__()
fourier_basis = torch.randn(in_c, out_c) * scale
self.register_buffer('_fourier_basis', fourier_basis)

def forward(self, x):
N, C, H, W = x.shape
x = rearrange(x, 'n c h w -> (n h w) c')
x = x @ self._fourier_basis
x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)

x = 2 * torch.pi * x
x = torch.sin(x)
return x

StyleGAN3 论文并没有讲为什么只用 sin,网上也很少有人讨论傅里叶特征的实现细节。我猜傅里叶特征并不是非得和傅里叶变换完全对应,毕竟它只是用来给神经网络提供更多信息,而没有什么严格的意义。只要把输入坐标分解成不同频率后,神经网络就能很好地学习了。

只用 sin 而不是同时用 sin, cos 后,似乎我们之前对 NTK 平移不变的推导完全失效了。但是,根据三角函数的周期性可知,只要是把输入映射到三角函数上后,网络主要是从位置间的相对关系学东西。绝对位置对网络来说没有那么重要,不同的绝对位置只是让所有三角函数差了一个相位而已。只用 sin 的神经网络似乎也对绝对位置不敏感。为了证明这一点,我把原来位于 [0, 1] 间的坐标做了一个幅度为 10 的平移。结果网络的误差几乎没变。

1
2
3
4
5
6
7
for epoch in tqdm(range(n_loops)):
x = fourier_feature(grid + 10)
output = model2(x)
loss = F.l1_loss(output, input_image)
optimizer.zero_grad()
loss.backward()
optimizer.step()

根据这些实验结果,我感觉是不是从 NTK 的角度来分析傅里叶特征完全没有必要?是不是只要从直觉上理解傅里叶特征的作用就行了?按我的理解,傅里叶特征在真正意义在于显式把网络对于不同频率的关注度建模出来,从而辅助网络学习高频细节。

总结

在这篇博文中,我们学习了傅里叶特征及其应用,并顺带了解其背后有关核回归、NTK 的有关理论知识。这些知识很杂乱,我来按逻辑顺序把它们整理一下。

为了解释为什么 NeRF 中的位置编码有效,傅里叶特征论文研究了用 MLP 拟合连续数据这一类任务中如何让 MLP 更好地学到高频信息。论文有两大主要结论:

  • 通过从 NTK 理论的分析,位置编码其实是一种特殊的傅里叶特征。这种特征具有平移不变性。因此,神经网络就像是在对某个输入信号做卷积。而我们可以通过调整傅里叶特征的参数来调整卷积的带宽,也就是调整网络对于不同频率的关注程度,从而使得网络不会忽略高频信息。
  • 傅里叶特征的频率不需要密集采样,只需要从任意一个分布随机稀疏采样。影响效果的关键是采样分布的标准差,它决定了傅里叶特征的带宽,也就决定了网络是否能关注到高频信息。

由于这些结论比较抽象,我们可以通过一个简单的二维图像拟合实验来验证论文的结论。实验表明直接将坐标输入给 MLP 不太行,必须将输入转换成傅里叶特征才能有效让网络学到高频信息。这个傅里叶特征可以是随机、稀疏的。

除了过拟合连续数据外,傅里叶特征的另一个作用是直接表示带宽有限信号,以实现在空域上的连续采样。StyleGAN3 在用傅里叶特征时,允许对输入坐标进行仿射变换,并且计算特征时只用了正弦函数而不是同时用正弦、余弦函数。这表明有关 NTK 的理论分析可能是没有必要的,主要说明问题的还是实验结果。

傅里叶特征论文仅研究了拟合连续数据这一类问题,没有讨论 Transformer 中位置编码的作用。论文中的一些结论可能无法适用。比如在大模型的位置编码中,我们还是得用密集的 sin, cos 变换来表示位置编码。不过,我们可以依然借助该论文中提到的理论分析工具,来尝试分析所有位置编码的行为。

只通过文字理解可能还不太够,欢迎大家尝试我为这篇博客写的 Notebook,通过动手做实验来加深理解。 https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/FourierFeature

推荐参考资料:

用代码示例理解核回归:https://towardsdatascience.com/kernel-regression-made-easy-to-understand-86caf2d2b844

直观理解 NTK: https://www.inference.vc/neural-tangent-kernels-some-intuition-for-kernel-gradient-descent/

傅里叶特征官方示例 PyTorch 实现: https://github.com/ndahlquist/pytorch-fourier-feature-networks/tree/master

论文速览

平时我自己读新论文时,往往简单看一下摘要和图表就能差不多明白文章的大意。而要用一篇博文来详细介绍新论文的设计动机、实现方法,需要花多得多的时间来写作。为了能分享更多的新论文,并且让我自己能够更好地整理论文间的关系,从这篇文章开始,我会不定期发表「论文速览」类文章,以简明扼要的文字来帮助计算机视觉的同行们快速了解新论文。

和发文频繁的自媒体相比,我的论文速览文章几乎不会展示论文的结果,仅会从科研角度介绍文章的贡献及方法,并给出我的批判性分析。和我之前的论文详解文章相比,这个新系列的文章会用更少的文字介绍背景,几乎不介绍方法细节,力求用更少的文字来表达关键信息,提升文字沟通效率。不过另一方面,我会花更多心思在介绍新论文与过往论文的关系上,为不熟悉某些背景知识的读者提供学习途径,以填补背景介绍上的空缺。

Diffusion Forcing

视频扩散模型普遍存在视频质量随着视频长度增加不断下降的问题。为此,论文作者提出了 Diffusion Forcing 这一建模任何序列生成问题的新范式:在训练该扩散模型时,序列中的每个元素会独立地添加不同强度的噪声。作者在简单视频生成、决策 (decision-making) 任务上验证了这种生成范式的有效性。我将主要从视频生成的角度介绍本工作。

以往工作

稍后我们会了解到,Diffusion Forcing 与此前两种主流序列生成范式密切相关:自回归生成与全序列扩散模型。

在序列的自回归生成中,模型会不断根据序列前 $n-1$ 个元素,预测下一个元素(第 $n$ 个元素)。自回归生成在 NLP 中最为常见,RNN、Transformer 用的都是这种生成方法。

扩散模型可以直接生成任何形状的数据。如果我们不把视频视为一种由图像组成的序列数据,而是将其视为一种「三维图像」,那么我们可以直接将 2D 图像扩散模型简单地拓展成 3D 视频扩散模型。这种做法在这篇论文中被称为「全序列扩散模型」。使用这一方法的早期工作有 DDPM 的作者提出的 Video Diffusion Models。Stable Diffusion 的作者也基于 LDM 提出类似的 Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models(Video LDM)工作。

全序列扩散模型仅能生成固定长度的视频。为了将其拓展到长视频生成,还是得将其和自回归结合起来。但是,自回归视频生成存在着生成质量与训练集质量不匹配的问题。Stable Video Diffusion 等工作参考 Cascaded Diffusion Models 的做法,通过给约束图像/视频帧加噪声缓解此问题。

AR-Diffusion: Auto-Regressive Diffusion Model for Text Generation 进一步探讨了自回归生成与全序列扩散模型的结合方法:在生成文本时,不同时刻的文本噪声不同,越早的文本上的噪声越少。无独有偶,FIFO-Diffusion: Generating Infinite Videos from Text without Training 提到了如何在预训练视频扩散模型上,以不同的噪声强度来生成不同的视频帧。或许是受到这些工作的启发,Diffusion Forcing 系统探讨了如何在训练时独立地给序列元素添加噪声。

科研动机

本工作的作者发现了自回归与全序列扩散模型的不足之处,并认为这两种生成范式是互补的:

  • 自回归不能在推理时加入新的优化目标,且存在着前述的训练采样质量不匹配导致的降质问题。
  • 全序列扩散模型不能生成变长的序列。

反过来讲:

  • 全序列扩散模型可以在推理加入 Classifier-guidance,且在固定序列长度内几乎不存在降质。
  • 自回归可以生成变长序列。

那能不能把二者结合起来呢?在推理时,我们希望序列是按时间顺序先后生成的。同时,从噪声强度这一维度上,我们希望每个元素能够从完全带噪到完全清晰逐渐生成。结合二者的方法是:在采样时,较早的元素有较少的噪声,较新的元素有较多的噪声,不同噪声强度的元素在序列中同时生成。比如采样分3步完成,共生成3帧,那么一步去噪会让$[x_1^{1/3\cdot T}, x_2^{2/3\cdot T}, x_3^{T}]$变成$[x_1^0, x_2^{1/3\cdot T}, x_3^{2/3\cdot T}]$。

为了实现这样的采样,我们在训练时就要让每一元素能够在其他元素和自己的噪声强度不同时也能顺利去噪。和以往的工作不同,本工作的作者发现,训练时我们不必和推理时一样固定每一元素的噪声强度,而是可以独立地对每帧随机采样噪声强度。

简单视频生成模型

本文的想法十分简洁,只需要在 DDPM 序列生成模型上把训练和推理时的噪声强度换掉即可。为了进一步了解本工作的方法,我们来看一下论文中有关视频模型的方法与实验。

整体上看,训练时,方法和 $\epsilon$ 预测的普通 DDPM 一样,只是不同帧噪声强度不一样。

在帧间关系上,Diffusion Forcing 用因果(causal)关系建模,即当前帧只能看到之前帧的信息。

具体来说,本工作用 RNN (准确来说是 GRU)的隐变量来建模之前帧传过来的信息。加入了 RNN 后,论文把本来很简单的公式变得无比复杂,不建议读者深究论文中有关 RNN 的内容。

在采样时,由于不同帧的噪声强度不同,现在我们需要定义一个二维的噪声强度表,表示每一帧在不同位置及不同去噪时刻的噪声强度。为了让每一时刻的噪声强度不同,一开始较新帧的去噪时刻会停留在原地。作者在附录中介绍了同时刻去噪的详细设计。

作者发现,Diffusion Forcing 的这种序列生成方式可以自然地推广到无限长度的视频生成上:在生成下一个片段时,不用滑动窗口,直接把 RNN 的初始隐变量初始化为上一个片段的隐变量输出。

作者用同样的 RNN 结构训练了两个基准模型:一个自回归模型,一个全序列因果扩散模型。定性结果表明,不管是在预定视频长度内,还是超出原本长度的长视频生成,Diffusion Forcing 的结果均好于基准方法。结果可以在官方项目网站上查看:

https://boyuan.space/diffusion-forcing/

批判性分析

论文开头说自回归模型在推理时缺少添加约束的方法。但这对视频生成来说并不致命,因为一般可以在训练时加入约束,推理时用 Classifier-free Guidance 就行了。

作者说出于简洁,他们在论文中用 RNN 实现了 Diffusion Forcing。但很明显 3D U-Net 才应该是直观上最简单实用的实现方法,毕竟最早期的视频扩散模型就是拿 3D U-Net 做的。在官方仓库中,有本科生帮他们实现了一个 3D U-Net 加时间注意力的模型,比原来视频模型效果要好。

我认为本文的视频生成基准方法设置过低。对于自回归视频生成/图生成视频,现在大家都会参照 Cascaded Diffusion, 对约束图像加噪并把噪声强度做为额外约束传入当前帧的生成模型。这种设计和 Diffusion Forcing 原理相似。为了体现新方法的好处,有必要跟这个更加强大的基准自回归方法做对比。

作者对于全序列视频扩散模型的设计也很奇怪。全序列视频扩散模型的初衷就是把视频当成 3D 图像来看待,允许帧间两两交换信息,只保证预定长度内的视频是连贯的。作者现在用 RNN 实现了一个因果版全序列视频模型,这个模型的表现肯定是不如非因果版的。虽然作者说 Diffusion Forcing 在因果视频生成上总是比全序列扩散模型更加连贯,我很怀疑去掉了因果这个条件后 Diffusion Forcing 还能否比得过。

Diffusion Forcing 在视频生成的主要好处应该体现在超出预定长度的长视频生成。因此,哪怕在预定长度内比不过全序列扩散模型,也没有关系。作者应该比较结合自回归和全序列扩散模型的方法,比如用 Stable Video Diffusion 这种图生视频模型,把上一个视频的末帧当作下一个视频的首帧约束,证明 Diffusion Forcing 在长视频生成上的优越性。

综上,我认为作者在视频生成任务上的实验是不够充分的。也的确,这篇论文有一半篇幅是在决策任务上,没有只讲视频生成任务。我相信 Diffusion Forcing 的设计会在长视频生成上缓解降质问题,这需要后续大公司的工作来跟进。但是,长视频的根本问题是记忆缺失,这一本质问题是 Diffusion Forcing 这种方法难以做好的。

这篇工作对我最大的启发是,我们一直把视频当成完整的 3D 数据来看待,却忘了视频可以被看成是图像序列。如果把视频当成 3D 数据的话,不同帧只能通过时序注意力看到其他帧在当前去噪时刻的信息;而对于序列数据,我们可以在不同帧的依赖关系上做更多设计,比如这篇工作的不同去噪强度。我很早前就在构思一种依赖更加强的序列生成范式:我们可不可以把序列中其他元素的所有去噪时刻所有信息(包括中间去噪结果及去噪网络的中间变量)做为当前元素的约束呢?这种强约束序列模型可能对多视角生成、视频片段生成等任务的一致性有很大帮助。由于生成是约束于另一个去噪过程的,我们对此去噪过程做的任何编辑,都可以自然地传播到当前元素上。比如在视频生成中,如果整个视频约束于首帧的去噪过程,那么我们用任意基于扩散模型的图像编辑方法来编辑首帧,都可以自然地修改后续帧。当然,我只是提供一个大致的想法,暂时没有考虑细节,欢迎大家往这个方向思考。

有人肯定也会想到能否把 Diffusion Forcing 拓展到图像像素间关系上。我认为要实现训练是完全没有问题的,问题出在推理上:Diffusion Forcing 在推理时需要预定义不同元素间的去噪强度。对于视频这种有先后顺序的数据,我们很自然地可以让越早的帧噪声强度越低。但是,对于图像来说,如何定义不同像素在不同去噪步数时的去噪强度并不是一个易解的问题。

近期,有两个大型多模态模型于同期公布:一个是来自 Meta 的 Transfusion,另一个是来自 Show Lab 和字节跳动的 Show-o 。好巧不巧,二者都宣称自己的模型是几乎最早将多模态任务用一个 Transformer 完成的,不需要借助额外的文本编码器实现图像生成,同时结合了自回归生成和扩散模型。我很好奇这两篇工作究竟有多少创新,于是快速扫完了这两篇论文,并简单给大家分享一下它们的核心内容。

在这篇文章中,我会快速介绍两篇工作的核心模型架构与部分实验结果。由于我仅对视觉任务比较熟悉,对语言和多模态没有那么了解,我的分析将主要围绕视觉任务。

论文 Arxiv 链接:

Transfusion: https://arxiv.org/pdf/2408.11039

Show-o: https://arxiv.org/pdf/2408.12528

读前准备

在阅读这两篇新工作时,建议大家先熟悉以 Transformer 为代表的自回归生成、以 DDPM、LDM、DiT 为代表的扩散模型、以 MaskGIT (Masked Generative Image Transformer), MAR (Masked autoregressive models, 于 Autoregressive Image Generation without Vector Quantization 论文中提出) 为代表的掩码自回归图像生成这三类生成模型,并简单了解此前较为先进的 Chameleon (Chameleon: Mixed-Modal Early-Fusion Foundation
Models
) 多模态模型。本文不会对这些知识做深入回顾,如果读者遇到了不懂的旧概念,请先回顾有关论文后再来看这两篇新文章。

自回归模型

自回归模型用于生成形如 $\mathbf{x}=\{x_1, x_2, …, x_n\}$ 这样的有序序列。自回归算法会逐个生成序列中的元素。假设正在生成第 $i$ 个元素,则算法 $F$ 会参考之前所有的信息 $\{x_j \mid j < i\}$,得到 $x_i = F(\{x_j \mid j < i\})$。比如:

自回归任务最常见的应用场景是文本生成。给定第一个词,生成第二个词;给定前两个词,生成第三个词……。

为了训练实现这一算法的模型,一般我们需要假设每个元素的取值是有限的。比如我们要建模一个生成单词的模型,每个元素是小写字母,那么元素的取值只有 a, b, c, ..., z。满足这个假设后,我们就可以像分类任务一样,用神经网络模型预测的类别分布建模已知之前所有元素时,下一个元素的分布,并用交叉熵损失函数去优化模型。这种训练任务被称为下一个词元预测 (next token prediction, NTP)。

用自回归生成建模某类数据时,最重要的是定义好每个元素的先后顺序。对于文本数据,我们只需要把每个词元 (token) 按它们在句子里的先后顺序编号即可。而对于图像中的像素,则有多种编号方式了。最简单的一种方式是按从左到右、从上到下的顺序给像素编号。

掩码自回归模型

由于图像的像素数很多,用自回归模型一个一个去生成像素是很慢的;另外,按从左到右、从上到下的顺序给像素编号显然不会是最合理的。为了提升自回归模型的速度和表现,研究者提出了掩码自回归模型。它做了两个改进:

1) 相比按序号一个一个生成元素的经典自回归模型,这种模型在每轮生成时可以生成多个像素(下图中的橙色像素)。

2) 相比从左到右、从上到下的固定顺序,像素的先后顺序完全随机(下图中的 (b) 和 (c) )。

由于这种方式下必须一次给模型输入所有像素,并用掩码剔除未使用的像素,所以这种自回归被叫做掩码自回归。

扩散模型

扩散模型将图像生成表示成一个噪声图像 $\mathbf{x}_T$ 从时刻 $T$ 开始随时间变化 $\mathbf{x}_t = F(\mathbf{x}_{t+1})$,最后得到目标图像 $\mathbf{x}_0$ 的过程。和输入输出为部分像素的自回归模型不同,扩散模型的输入输出自始至终都是完整图像。

为了减少扩散模型的计算量,参考 Latent Diffusion Model (LDM) 的做法,我们一般会先用一个自编码器压缩图像,再用扩散模型生成压缩过的小图像。正如 NLP 中将文本拆成单词、标点的「词元化」(tokenize) 操作一样,这一步操作可以被称为「图块化」(patchify)。当然,有些时候大家也会把图块叫做词元,把图块化叫做图像词元化。

严格来说,本文讲到的「像素」其实是代表一个图块的图像词元。用「像素」是为了强调图像元素的二维空间信息,用「图像词元」是强调图像元素在自回归模型中是以一维序列的形式处理的。

有人认为,掩码自回归模型是一种逐渐把纯掩码图像变成有意义图像的模型,它和逐渐把纯噪声图像变成有意义图像的扩散模型原理类似。因此,他们把掩码自回归模型称为离散扩散模型。还有人认为扩散模型也算一种更合理的自回归,每轮输入一个高噪声图像,输出一个噪声更少的图像。但这些观点仅仅是从称呼上统一两种模型,两种模型在实现上还是有不少差别的。

Chameleon

Chameleon 似乎是此前最为先进的多模态模型,它是这两篇新工作的主要比较对象。在语言模型的基础上,Chameleon 并没有对图像的处理多加设计,只是以离散自编码器(如 VQGAN)的编码器为图像词元化工具,以其解码器为图像词元化还原工具,让被词元化后的图像词元以同样的方式与文本词元混在一起处理。

功能与效果

看方法前,我们先明确一下两个多模态模型能做的任务,及各任务的输入输出。

Transfusion 是一个标准多模态模型,也就是一个输入输出可以有图像词元的语言模型。它输入已知文本和图像,输出后续文本和图像。

基于这个多模态模型,可以做文生图任务。

这个模型似乎没有为特定任务设置特殊词元,所有图像功能完全靠文本指定。因此,要做图像编辑任务的话,需要在带文本标注的图像编辑数据集上微调。文章指出,只需要在一个仅有 8000 项数据的数据集上微调就能让模型具有一定的编辑能力。

相比之下,Show-o 可以在序列前多输入一个区分任务的特殊词元。因此,Show-o 可以完成多模态理解(输入多模态,输出文本描述)、图像生成(根据文字生成图像或填补图像)、多模态生成(输入输出都包含图片、文本)等丰富的任务。似乎特殊词元仅有多模态理解 (MMU, Multi-modal Understanding MMU)、文生图 (T2I, Text to Image) 这两种。

Transfusion 的基础模型在微调后才能做根据文本提示来编辑图像的任务,而 Show-o 的基础模型默认是在此类带有文本提示的图像编辑数据集上微调的。

方法

对于熟悉此前图像生成模型、语言模型的研究者来说,这两篇工作都仅用了现有技术,核心方法非常易懂。这两篇工作并不是试图开发一种新的图像生成技术,而是在考虑如何更好地将现有图像模型融入多模态模型。

在读新技术之前,我们先以 Chameleon 为例,看一下之前的多模态模型是怎么做生成的。在我看来,之前的多模态模型不应该叫「多模态模型」,而应该叫「强行把图像当成词元的语言模型」。语言模型在处理文本时,文本中的词语、标点会根据规则被拆成「词元」,成为模型的最小处理单位。然而,要用同样的方式处理图像,就要定义什么是「图像词元」。受到之前图像离散压缩模型(以 VQGAN 为代表)的启发,大家用一个编码器把图像切成图块,每个图块可以用一个 1, 2, 3 这样的整数序号表示,再用一个解码器把用序号表示的图块翻译回真实图像。这里的带序号图块就和文本里的单词一样,可以用「词元」来表示。文本、图像都被统一成词元后,就能用标准 Transformer 的下一个词元预测任务来训练多模态模型。

如下图所示,训练时,文本基于程序规则变成词元,而图像经过一个编码器模型变成词元。两类词元被当成一类数据,以同样的方式按下一个词元预测任务训练。生成时,多模态模型自回归地生成所有词元,随后文本词元基于程序规则恢复成文本,而图像词元通过解码器模型恢复成图像。

这种多模态模型最大的问题是没有充分设计图像词元生成,还是暴力地用 Transformer 自回归那一套。虽然有些模型会多加入一些图像生成上的设计,比如 LaVIT (Unified Language-Vision Pretraining in LLM with Dynamic Discrete Visual Tokenization) 用扩散模型来做图像词元的解码,但核心的图像词元生成还是离不开标准自回归。

Transfusion 和 Show-o 的设计初衷都是引入更先进的图像生成技术来改进图像词元生成。先看 Show-o。要改进标准的一个一个按顺序生成图像词元的图像自回归模型,最容易想到的做法就是按照 MaskGIT, MAR 那一套,将标准自回归换成掩码自回归。在做掩码自回归时,像素的先后顺序完全随机,且一次可以生成多个像素。另外,图像词元之间可以两两互相做交叉注意力,而不用像文本词元一样只能后面的对前面的做交叉注意力。

Show-o 莫名其妙地把自己的图像生成模型称为离散扩散模型。如前文所述,叫离散扩散模型还是掩码自回归,只是一个称呼上的问题。由于问题建模上的重大差异,大家一般还是会把扩散模型和掩码自回归看成两类模型。

而 Transfusion 更加激进地在革新了多模态模型中的图像生成方式。现在最好的图像生成技术不是扩散模型吗?我们干脆直接把整个扩散模型搬过来。于是,在 Transfusion 生成多模态内容时,程序会交替执行两种模式:在语言模型模式下,程序按标准自回归逐个生成文本词元。一旦生成了特殊词元 BOI (begin of image),就切换到扩散模式。在扩散模式下,程序按 DiT, SD 3 (Stable Diffusion 3) 那种标准扩散模型的方式,一次性生成所有图像词元。结束此模式后,程序往词元序列里填一个 EOI (end of image),重返语言模型模式。

同理,在训练时,两种模态也用不同的任务来训练。语言模型老老实实地按下一个词元预测训练,而扩散模型那部分就按照训练扩散模型的标准方式,先给所有图像词元加噪,再预测噪声。因此,只看图像生成任务的话,Transfusion 更像 SD 3 这种文生图模型,而不像此前的基于语言模型的多模态模型。

Transfusion 和 SD 3 之间的最大区别在于,文本词元还是按照语言模型那一套,只能在交叉注意力看到之前的文本词元。而图像词元之间两两都能看到。这种交叉注意力的设计和 Show-o 是一模一样的。当然,由于现在文本也会在同一个 Transformer 里处理,所以 Transfusion 自己就扮演了解读文本的工作,而不像 SD 3 那样还需要单独的文本编码器。

定量评测结果

我们最后来看一下两篇文章展示的定量评测结果。

Transfusion 用了许多篇幅来展示它与 Chameleon 之间的对比。从数值指标上看,Transfusion 全面领先 Chameleon。明明没有对文本任务做特别的优化,Transfusion 却在文本任务超越了 Chameleon,这挺令人惊讶的。为了探究这一现象的原因,作者从同一个预训练语言模型开始,以同样的配置训练了 Transfusion, Chameleon。结果显示,相比加入图像扩散模型,加入 Chameleon 那种离散图像词元对文本指标的损害更大。作者猜测,这是因为扩散模型更加高效,模型能够把更多精力放在文本任务上。

而从图像生成模型的对比上看,Transfusion 比之前多数文生图模型都要好,只是比顶尖文生图模型 SD 3 要差一点。

再来看 Show-o 的评测结果。Show-o 在部分文本指标上超过了之前的语言模型。作者也坦言,这些指标仅表明 Show-o 有潜力在文本任务上做到顶尖。

Show-o 也展示了图像任务的指标。和 Transfusion 一样,Show-o 展示了表示图像质量的 COCO FID 指标以及评价文本图像匹配度的 GenEval 指标。Show-o 在图像指标上超越了此前多数多模态模型,且超越了 Stable Diffusion 2.1 等图像生成模型。但是其图像指标比 Transfusion 还是差了不少。Show-o 的最大优点是需要的图像训练数据远远少于其他模型。

总结与讨论

此前多模态模型都只是强行把图像变成离散图像词元,再用标准自回归来生成图像词元。为了改进这些多模态模型,无独有偶,Transfusion 和 Show-o 都用到了更先进的图像生成技术。Show-o 将标准自回归改成了更强大的掩码自回归,而 Transfusion 激进地引入了完整的图像扩散模型,并把文本生成和图像生成当成两个相对独立的任务。二者的相同之处如下:

  • 两个多模态模型都用同一个 Transformer 来处理文本和图像。
  • 两个模型的 Transformer 都使用了同样的交叉注意力机制。文本词元只能看到此前的图像、文本词元,而图像词元可以看到此前的文本词元和当前所有图像词元。

二者的不同之处在于:

  • Transfusion 使用标准扩散模型实现图像生成,而 Show-o 使用掩码自回归实现图像生成。Show-o 强行将自己的图像生成模型称为「离散扩散模型」,有借扩散模型的名头宣传之嫌。
  • Transfusion 没有用特殊词元来区分不同任务。要用 Transfusion 编辑图像,需要在基础模型上用图像编辑数据集微调。Show-o 用特殊词元来区分不同任务,默认支持文本理解、图像生成、图像编辑、多模态生成等多种任务。
  • 二者在文本、图像指标上都超越了之前的多模态模型。但二者相互对比之下,Transfusion 的表现更好,而 Show-o 需要的训练资源少得多。

我再来谈一下我看完这两篇文章后的一些感想。此前我对多模态模型很不感兴趣,觉得它们就是在语言模型的基础上,强行加入图像信息,然后加数据、加显卡,大火乱炖,没有太大的创新。而 Transfusion 和 Show-o 的设计令我眼前一亮。我觉得这两篇文章的结果再次表明,图像和文本的处理方法不应该是一致的,不能强行把文本自回归那套方法直接搬到图像任务上。不管是换成 MaskGIT 掩码自回归那一套,还是完全用扩散模型,都比之前的方法要好。

而究竟是掩码自回归好还是扩散模型更好呢?在我看来,扩散模型是更好的。文本是离散的,而图像是连续的。这种连续性不仅体现在图像的颜色值上,还体现在图像像素间的位置关系上。强行用 Transformer 自回归生成图像,一下子把这两种连续信息都破坏了。而何恺明团队近期的 MAR 工作则试图找回图像在颜色值上的连续性,但依然无法充分利用空间上的连续性。相比之下,扩散模型每步迭代都是对整幅图像做操作,不会破坏图像的连续性。这两个多模态工作也反映了这一点,Transfusion 的表现要比 Show-o 好很多。

在生成图像时,Transfusion 的行为几乎和 SD 3, FLUX.1 这些文生图模型一样了。两类模型的区别在于,SD 3 它们得用一个预训练的语言处理模型。而 Transfusion 用同一个 Transformer 来处理文本、图像信息。尽管现在 Transfusion 还没有超过 SD 3,但我认为文生图任务本质上是一个多模态任务,这两类模型的后续发展路线很可能会交汇到一起。文生图也应该从多模态中汲取经验。比如我们在 SD 3 的基础上,加入一些语言任务上的学习,说不定能进一步提升文生图时的图文匹配度。

当然,仅根据我读完这两篇文章后有限的认知,我认为多模态并不是一个值得广大科研人员投入的研究方向,而只适合有足够资源的大公司来做。其根本原因是验证多模态设计的代价太大了。在图像生成领域,要验证一个生成模型好不好,我们拿 ImageNet,甚至拿只有几万张图像的人类数据集训练及评测,都很有说服力。而多模态任务必须在大量文本、图像数据集上训练,评测文本、图像上的多种指标,小一点的团队根本做不来。这样,各种小创新都难以验证。而没有各种各样的小创新,领域也就很难快速发展。所以多模态只能不断从纯语言生成、纯图像生成方向找灵感,很难有仅属于多模态任务的创新。