【深度学习】One Model to Learn Them All详解

Kaiser, Lukasz, et al. “One Model To Learn Them All.” arXiv preprint arXiv:1706.05137 (2017).

概述

Google于2017年6月16日在arxiv上提交了这篇论文，甫一问世立刻引发各方关注。除了标题劲爆之外，本文的野心和气魄令人惊叹，实验也确实给出了一些相当有信息量的结果。

项目的github页面给出了基于tensorflow的源码，完成度一般。本文结合此源码讲解系统结构。

系统

问题

本文尝试用一个通用模型解决跨领域的各类人工智能问题，例如：

- 图像分类（图像 -> 类标）

- 看图说话（图像 -> 自然语言）

- 翻译（自然语言 -> 自然语言）

- 语义分割（自然语言 -> 分割+类标）

各领域输入输出的信息类别不同，在本文中称为不同形态（modality）。

这种向着大一统模型的努力并非本文首创，创新点如下：

工作	领域	任务
以往文章	单一	多个
e.g.	翻译	英翻法+英翻德
本文	多个	多个
e.g.	翻译+图像分类	英翻法+英翻德+1000类分类

结构

为了适应不同形态的输入和输出，本文的网络被抽象成如下结构：

input

、

output

：相同/不同形态的数据（例如图像和类标）。

source

、

target

：系统内部的表达。

系统通过三个部分：

modality_in

,

modality_out

以及

body

来完成数据流。

三大理念

绝大部分计算量都集中在

body

网络中，两个

modality

网络设计尽量精简。

系统内部的表达（

target

,

source

）尺寸不固定。

对于相同形态的不同问题（例如“看图说话”和“英翻德”的输出），使用相同的

modality

网络

以下分别介绍

body

网络和

modality

网络的具体结构。

body网络-模块

body

网络各部分都由三种基本模块（block）构成，以下一一介绍。

卷积模块

Conv

子模块-

ConvStep

SepConv

：分层卷积，类似这篇博客介绍过的Factorized卷积。有三个超参数：扩展

d

，步长

s

，核尺寸

f

。

LayerNorm

: 分层归一化。

实现参见/models/common_layers.py中

conv_block_internal

函数/

subseparable_conv_block

函数。

Conv

构成

使用上述子模块组成卷积模块

Conv

：



主体结构是两个residual结构。最后虚线的dropout只在训练时使用。

实现参见/models/slicenet.py中

multi_conv_res

函数。

注意力模块

Attention

在处理时间序列信号时，以往工作多采用RNN, LSTM类型系统。此类系统每一个时刻输出取决于前时刻记忆，天然地不利于并行计算，在训练时尤其耗时。

本文则利用

Attention

模块，能够同时处理输入序列的各个元素。细节参考自同一团队的论文

Attention is All You Need

1。

该文对应模型为/models/transformer.py

子模块-

AddTiming

由于本文不再将时间序列顺序输入系统，所以需要额外告知系统每一元素在序列中的相对位置。

首先考虑1维信号x，输入尺寸为L×D。

对于D个通道，按照指数坐标均匀设置从2π到10000⋅2π的周期，共有D/2个采样:T0,T1...TD/2−1。



每通道的时间信号为上述波长的正弦/余弦，其自变量范围为[0,L−1]（e.g. D=12, L=128）：



得到的时间信号和输入信号尺寸相同，直接相加。

y=x+timing

对于n维信号，输入尺寸为L1,L2...Ln,D。

采样的周期数量只需要是1维情况的1/n：T0,T1...T(D/2n)−1。

对于每一维度，生成D/(2n)对不同频率的正弦/余弦信号，扩展为L1⋅L2⋅...Ln大小。

共有D个时间信号，分别加到D个通道上。

实现参见/models/common_attention.py中

add_timing_signal_1d

和

add_timing_signal_2d

函数。

子模块-

Dot-Prod Attention

注意力网络有三个输入

Q（Query）：想要考察的一组当前对象属性。尺寸为Lq×Dk。

K（Key）：已经存在的一组参考对象属性。尺寸为Lkv×Dk。

V（Value）：参考对象的值。尺寸为Lkv×Dv。

输出：当前对象的值。尺寸为Lq×Dv。如下图计算。



其物理意义是，考察Q和K中元素的两两相似程度，用相似程度作为权重，将V的加权和作为输出。

实现参见/models/common_attention.py中的

dot_product_attention

函数。

子模块-

Multi-Head Attention

首先在前述

Dot-Product Attention

的三个输入端添加线性投影；

之后将g个这样的结果串接起来；

最后重新投影成系统内部表达需要的维度。



实现参见/models/common_attention.py中的

multihead_attention

函数。

Attention构成

Attention模块有两个输入：尺寸为Lq×C的

query

，以及尺寸为Lm×C的

memory

。



首先为查询添加时间信息。

左侧的第一个

Multi-Head Attention

模块施加在输入的查询上，在其Lq个元素之间建立关联。

右侧的第二个

Multi-Head Attention

综合当前查询的Lq个元素和原有记忆Lm个元素之间的关系，输出Lq个查询结果。

实现参见/models/slicenet.py中的

attention

函数。

混合专家模块

Sparsely-Gated MOE

MOE类模块能够在不增加计算量的前提下，构造具有海量参数的模型，大幅提高模型表达能力。细节参看Google Brain团队的Outrageously Large Neural Networks2.

整个模块包含若干并行的“专家”Ei(x)。它们的结构相同，参数不同。都是重复若干层的线性网络+激活函数。

实现参见/utils/expert_utils.py中

FeedForwardExpert

函数。

另外有一个和输入有关的门函数G。其中G(x)是一个系数的n维向量，如果G(x)i=0，则不必计算Ei(x)。

本文采用的门函数如下：



其中添加的噪声

StandardNormal

相当于一个平滑项，其强度由Wnoise控制。

实现参见/utils/expert_utils.py中

NoisyTopKGating

函数。

输出由所有专家子模块通过门函数加权得到：

y=∑i=1nG(x)iEi(x)

其中n=280/60,k=4。

实现参见/utils/expert_utils.py中

DistributedMixtureOfExperts

函数。

body网络-构成

body

网络由如下三部分构成：



系统的输入和输出都是时间序列（非时间信号可以看做长度为1的特例）。

Encoder

部分处理将

source

编码；

Mixer

部分将编码后的

source

和系统此时刻之前的记忆综合起来，生成编码后的

memory

。

Decoder

部分从编码后的源和记忆生成

target

表达。

Input Encoder

源信号

source

首先添加时间信息，通过3次卷积，并通过一个MOE模块。

之后，与自身重复进行3次

Attention

，相当于充分关联输入序列。

最后得到编码后的源信息

encoded source

。



实现参见/models/multimodel.py，

MultiModel

类

model_fn_body_sharded

函数84-106行。源码和论文无法一一对应。

I/O Mixer

首先将编码后的源信息和记忆信息通过

Attention

进行混合。

之后经过concat操作压缩一维。（此处不详）

最后将混合信息通过与

Encoder

类似的2次自身

Attention

操作，获得编码后的记忆

encoded memory

。



实现参见/models/multimodel.py，

MultiModel

类

model_fn_body_sharded

函数119-142行。源码和论文无法一一对应。

Decoder

首先将编码后的源信息和记忆信息串接起来。

而后经过两个卷积模块。

最后通过与

Encoder

类似的4次自身

Attention

操作获得目标信息

target

。



实现参见/models/multimodel.py，

MultiModel

类

model_fn_body_sharded

函数108-116行。

具体实现：/models/slicenet.py中

slicenet_middle

函数。源码和论文无法一一对应。

modality网络

对于同一形态的信息，modality网络有三种作用：

in

网络：把原始输入

input

转化为源信息

source

out

网络：把目标信息

target

转化为输出

output

regress

网络：把前时刻的目标信息转化成记忆

memory

论文中没有提到

regress

网络，直接用

out

+

in

代替，但在源码中有所体现。

基类实现参见/utils/modality.py中的

Modality

类。

bottom

,

targets_bottom

,

top

函数分别对应上述三个网络。

不同形态信息的具体实现参见/models/modalities.py

实验与总结

略去细节，只说值得注意的现象和结论。

在没有仔细fune-tune的前提下，本文结论只比state-of-art稍逊色。

同时在不同领域的多个任务上训练，几乎不会损害单个任务的精度。

对于小数据集任务，同时训练其他任务甚至能够提升本任务的表现。即使是毫不相关形态之下的问题。

传统上用于某种形态问题的模块（例如用于语言的attention机制和MOE）能够对其他形态的问题有所帮助。

总体来说，本文在大一统模型的道路上又前进了一步。反观本文的三大设计理念，会发现其更接近人的行为方式：

本文	人类	实例
绝大部分计算量都集中在 body 网络中， modality 网络设计尽量精简。	复杂的思维组件负责处理不考虑形态的抽象概念；简单的输入输出组件负责处理和表达不同形态的具体信号	脑补很强大，眼耳口鼻很粗糙
系统内部的表达尺寸相同，但不固定。	不同复杂程度的抽象概念使用不同长度的信息量来存储。	越常用的概念表达越简单
对于相同形态的不同问题，使用相同的 modality 网络	同类的不同任务使用相同的输入输出组件	用同样的耳朵听不同的语言

本文的源码部分还不完善，有待观望。

Vaswani, Ashish, et al. “Attention Is All You Need.” arXiv preprint arXiv:1706.03762 (2017). ↩
Shazeer, Noam, et al. “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.” arXiv preprint arXiv:1701.06538 (2017). ↩