MMoE是一种多任务学习模型，它改进了传统的SharedBottom结构，通过引入多个专家网络（Expert）和门控机制（Gate），实现了根据不同任务动态调整专家网络组合权重的能力。每个任务拥有独立的Gate，允许网络参数根据输入和任务类型自适应变化。在代码实现中，MMoE创建了多个Expert和Gate网络，通过前向传播计算每个任务的专家输出，形成任务特定的特征表示。
一、MoE简介
MMoE是在MoE的基础之上，引入多任务而来，关于MoE的相关解读，可以参考我的另一篇博客：【论文阅读】Adaptive Mixtures of Local Experts
二、MMoE简介
关于MMoE的论文阅读在我的另一篇博客中已经讲到，具体可以参考：【推荐算法论文阅读】Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts
MMoE的模型结构如下图（c）所示：
MMoE的创新之处在于跳出了Shared Bottom那种将整个隐藏层一股脑的共享的思维定式，而是将共享层有意识的（按照数据领域之类的）划分成了多个Expert，并引入了gate机制，得以个性化组合使用共享层。
MoE共享层将大的Shared Bottom网络拆分成了多个小的Expert网络（如图所示，拆成了三个，并且保持参数个数不变)。我们把第i个Expert网络的运算记为
f
i
(
x
)
f_i(x)
fi(x)，然后Gate操作记为
g
(
x
)
g(x)
g(x)，他是一个n元的softmax值（n是Expert的个数，有几个Expert，就有几元），之后就是常见的每个Expert输出的加权求和，假设MoE的输出为y，那么可以表示为：
y
k
=
h
k
(
f
k
(
x
)
)
y_{k}=h^{k}\left(f^{k}(x)\right)
yk=hk(fk(x))
k是多任务中task的个数。
如果只是这样的话，要完成多任务还得像Shared Bottom那样再外接不同的输出层，这样一搞似乎这个MoE层对多任务来说就没什么用了，因为它无法根据不同的任务来调整各个Expert的组合权重。所以论文的作者搞了多个Gate，每个任务使用自己独立的Gate，这样便从根源上，实现了网络参数会因为输入以及任务的不同都产生影响。于是，我们将上面MoE输出稍微改一下，用
g
k
(
x
)
g^k(x)
gk(x)表示第k个任务的门就得到了MMoE的输出表达：
f
k
(
x
)
=
∑
i
=
1
n
g
k
(
x
)
i
f
i
(
x
)
g
k
(
x
)
=
softmax
(
W
g
k
x
)
f^{k}(x)=\sum_{i=1}^{n} g^{k}(x)_{i} f_{i}(x) \\g^{k}(x)=\operatorname{softmax}\left(W_{g k} x\right)
fk(x)=i=1∑ngk(x)ifi(x)gk(x)=softmax(Wgkx)
Gate：把输入通过一个线性变换映射到nums_expert维，再算个softmax得到每个Expert的权重
Expert：简单的几层全连接网络，relu激活，每个Expert独立权重
三、MMoE代码实现
MMoE 借鉴 MoE 的思路, 引入多个 Experts (即多个 NN 网络) 网络, 然后再对每个 task 分别引入一个 gating network, gating 网络针对各自的 task 学习 experts 网络的不同组合模式, 即对 experts 网络的输出进行自适应加权. 说实话, 这一点非常像 Attention, Experts 网络学习出 embedding 序列, 而 gating 网络学习自适应的权重并对 Experts 网络的输出进行加权求和, 得到对应的结果之后再分别输入到各个 task 对应的 tower 网络中. 注意 gating 网络的数量和任务的数量是一致的。
3.1 初始化创建 Experts 和 Gate 网络
class MMoE(Layer):
"""
Multi-gate Mixture-of-Experts model.
"""
def __init__(self,
units, ## 隐藏层单元个数
num_experts, ## Experts 的个数，也就是N
num_tasks, ## 任务个数，也就是K
use_expert_bias=True,
use_gate_bias=True,
expert_activation='relu',
gate_activation='softmax',
## .... 其他参数
**kwargs):
def build(self, input_shape):
"""
这里我们假设输入 tensor 的 shape 为 [B, I]
其中 B 为 Batch_size, I 为Input Embedding 的大小
隐藏层 units 单元个数使用 E 表示
Experts 网络的个数设为 N
Task 任务的个数设置为 K
"""
assert input_shape is not None and len(input_shape) >= 2
input_dimension = input_shape[-1] # I
"""
初始化 Experts 网络, 其大小为 [I, E, N],
其中 I 为输入 embedding 的大小, E 为 Experts 网络的输出结果大小,
N 为 Experts 网络的个数
"""
self.expert_kernels = self.add_weight(
name='expert_kernel',
shape=(input_dimension, self.units, self.num_experts),
initializer=self.expert_kernel_initializer,
regularizer=self.expert_kernel_regularizer,
constraint=self.expert_kernel_constraint,
)
"""
初始化 Experts 网络的 Bias, 大小为 [E, N]
"""
if self.use_expert_bias:
self.expert_bias = self.add_weight(
name='expert_bias',
shape=(self.units, self.num_experts),
initializer=self.expert_bias_initializer,
regularizer=self.expert_bias_regularizer,
constraint=self.expert_bias_constraint,
)
"""
初始化 Gate 网络, 注意 Gate 网络的个数和 Task 的个数相同, 均为 K,
因此 self.gate_kernels 列表的大小为 K, 每个 Gate 中 weight 的
大小均为 [I, N], I 为输入 Embedding 的大小, 而 N 为 Experts 网络的个数
Gate 网络的输出结果保存着各 Experts 网络的权重系数
"""
self.gate_kernels = [self.add_weight(
name='gate_kernel_task_{}'.format(i),
shape=(input_dimension, self.num_experts),
initializer=self.gate_kernel_initializer,
regularizer=self.gate_kernel_regularizer,
constraint=self.gate_kernel_constraint
) for i in range(self.num_tasks)]
"""
初始化 Gate 网络的 Bias, self.gate_bias 大小为 K,
每个 Bias 的大小为 (N,)
"""
if self.use_gate_bias:
self.gate_bias = [self.add_weight(
name='gate_bias_task_{}'.format(i),
shape=(self.num_experts,),
initializer=self.gate_bias_initializer,
regularizer=self.gate_bias_regularizer,
constraint=self.gate_bias_constraint
) for i in range(self.num_tasks)]
self.input_spec = InputSpec(min_ndim=2, axes={-1: input_dimension})
super(MMoE, self).build(input_shape)
3.2 MMoE 网络的前向传播
def call(self, inputs, **kwargs):
gate_outputs = []
final_outputs = []
# f_{i}(x) = activation(W_{i} * x + b), where activation is ReLU according to the paper
"""
inputs 输入 Tensor 的大小为 [B, I],
self.expert_kernels 的大小为 [I, E, N],
其中 I 为输入 embedding 大小, E 为 Experts 网络的输出大小, N 为 Experts 的个数
tf.tensordot(a, b, axes=1) 相当于 tf.tensordot(a, b, axes=[[1],[0]]),
因此 expert_outputs 的大小为 [B, E, N]
"""
expert_outputs = K.tf.tensordot(a=inputs, b=self.expert_kernels, axes=1)
# Add the bias term to the expert weights if necessary
if self.use_expert_bias:
expert_outputs = K.bias_add(x=expert_outputs, bias=self.expert_bias)
"""
加上 Bias 以及通过激活函数 (relu) 后, expert_outputs 大小仍为 [B, E, N]
"""
expert_outputs = self.expert_activation(expert_outputs)
# g^{k}(x) = activation(W_{gk} * x + b), where activation is softmax according to the paper
"""
针对 K 个 Task 分别学习各自的 Gate 网络, 这里采用 for 循环实现,
其中 inputs 的大小为 [B, I],
gate_kernel 的大小为 [I, N], 其中 I 为输入 embedding 的大小,
而 N 为 Experts 的个数. 因此 K.dot 对 inputs 和 gate_kernel 进行矩阵乘法,
得到 gate_output 的大小为 [B, N].
注意 gate_activation 为 softmax, 因此经过 Bias 以及 gate_activation 后,
gate_output 的大小为 [B, N], 保存着各 Experts 网络的权重系数
"""
for index, gate_kernel in enumerate(self.gate_kernels):
gate_output = K.dot(x=inputs, y=gate_kernel)
# Add the bias term to the gate weights if necessary
if self.use_gate_bias:
gate_output = K.bias_add(x=gate_output, bias=self.gate_bias[index])
gate_output = self.gate_activation(gate_output)
gate_outputs.append(gate_output)
# f^{k}(x) = sum_{i=1}^{n}(g^{k}(x)_{i} * f_{i}(x))
"""
gate_outputs 为大小等于 K (任务个数) 的列表, 其中 gate_output 的大小等于 [B, N],
而 expert_outputs 的大小为 [B, E, N];
因此, 首先对 gate_output 使用 expand_dims, 按照 axis=1 进行, 得到
expanded_gate_output 大小为 [B, 1, N];
K.repeat_elements 将 expanded_gate_output 扩展为 [B, E, N],
之后再乘上 expert_outputs, 得到 weighted_expert_output 大小为 [B, E, N];
此时每个 Experts 网络都乘上了对应的系数, 最后只需要对各个 Experts 网络的输出进行加权
求和即可, 因此 K.sum(weighted_expert_output, axis=2) 的结果大小为 [B, E];
"""
for gate_output in gate_outputs:
expanded_gate_output = K.expand_dims(gate_output, axis=1) ## [B, 1, N]
weighted_expert_output = expert_outputs * K.repeat_elements(expanded_gate_output, self.units, axis=1) ## [B, E, N]
final_outputs.append(K.sum(weighted_expert_output, axis=2)) ## [B, E]
return final_outputs
395
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
