【深度学习】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 Need1。

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

子模块-AddTiming

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

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

对于D个通道 ，按照指数坐标均匀设置从2\pi到10000\cdot 2\pi的周期，共有D/2个采样:T_0, T_1...T_{D/2-1}。

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

得到的时间信号和输入信号尺寸相同，直接相加。
y=x+timing

对于n维信号 ，输入尺寸为L_1,L_2...L_n,D。
采样的周期数量只需要是1维情况的1/n：T_0, T_1...T_{(D/2n)-1}。
对于每一维度 ，生成D/(2n)对不同频率的正弦/余弦信号，扩展为L_1\cdot L_2\cdot ...L_n大小。
共有D个时间信号，分别加到D个通道 上。

实现参见/models/common_attention.py中add_timing_signal_1d和add_timing_signal_2d函数。

子模块-Dot-Prod Attention

注意力网络有三个输入
Q（Query）：想要考察的一组当前对象属性。尺寸为L_q \times D_k。
K（Key）：已经存在的一组参考对象属性。尺寸为L_{kv} \times D_k。
V（Value）：参考对象的值。尺寸为L_{kv} \times D_v。

输出：当前对象的值。尺寸为L_q \times D_v。如下图计算。

其物理意义是，考察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模块有两个输入：尺寸为L_q\times C的query，以及尺寸为L_m\times C的memory。

首先为查询添加时间信息。
左侧的第一个Multi-Head Attention模块施加在输入的查询上，在其L_q个元素之间建立关联。
右侧的第二个Multi-Head Attention综合当前查询的L_q个元素和原有记忆L_m个元素之间的关系，输出L_q个查询结果。

实现参见/models/slicenet.py中的attention函数。

混合专家模块Sparsely-Gated MOE

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

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

实现参见/utils/expert_utils.py中FeedForwardExpert函数。

另外有一个和输入有关的门函数G。其中G(x)是一个系数的n维向量，如果G(x)_i=0，则不必计算E_i(x)。
本文采用的门函数如下：

其中添加的噪声StandardNormal相当于一个平滑项，其强度由W_{noise}控制。

实现参见/utils/expert_utils.py中NoisyTopKGating函数。

输出由所有专家子模块通过门函数加权得到：
y=\sum_{i=1}^n G(x)_iE_i(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网络	同类的不同任务使用相同的输入输出组件	用同样的耳朵听不同的语言

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

---

1. Vaswani, Ashish, et al. “Attention Is All You Need.” arXiv preprint arXiv:1706.03762 (2017). ↩︎

2. Shazeer, Noam, et al. “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.” arXiv preprint arXiv:1701.06538 (2017). ↩︎