Zhangzhe's Blog

The projection of my life.

0%

Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts

URL

https://dl.acm.org/doi/pdf/10.1145/3219819.3220007

TL;DR

  • 本文提出一种多任务神经网络结构,称为 Multi-gate Mixture-of-Experts 简称 MMOE
  • 与传统多任务网络共享 bottom 相比,该结构可以在 任务相关性较弱 的情况下有较好的鲁棒性
  • MMOE 中的 Multi-gate 本质就是一种 Softmax Attention,针对不同的任务给出不同的专家组合

Algorithm

总体网络结构

mmoe1.png

其中 AB 分别表示两个 任务Tower 表示与任务相关的金字塔头部,Expert 表示专家 bottom network,Gate 表示任务相关的权重生成网络

  • 上图 a 表示 Bottom Shared Architecture: 传统共享 bottom network 的多任务结构,共享 bottom network (Backbone) 进行特征提取,将提取到的特征分别送入任务相关头部,缺点是如果任务 AB 之间的相关性较弱,那么共用一个 Backbone 是危险的

  • 上图 b 表示 One-gate Mixture-of-Experts Architecture,有多个 ExpertBackbone 进行不同维度的特征提取,只有一个 Gate network 用于给每个任务生成 Expert 权重

  • 上图 c 表示 Multi-gate Mixture-of-Experts Architecture,有多个 ExpertBackbone 进行不同维度的特征提取,每个任务有单独的 Gate network 生成唯一的 Expert 权重

数学定义

  • Bottom Shared

    yk=hk(f(x))y_k = h^k(f(x))

    其中 kk 表示 任务数f()f() 表示 shared bottomhkh^k 表示 第 k 个任务的 Tower

  • OMOE

    yk=hk(i=1ng(x)ifi(x))y_k=h^k(\sum_{i=1}^n g(x)_i f_i(x)), where i=1ng(x)i=1,  g(x)Rn\sum_{i=1}^n g(x)_i=1,\ \ g(x) \in \mathbb R^n

    其中 nn 表示 专家数gg 表示 Gate network(由于 g(x)g(x) 要经过 Softmax,使得 logits -> prob,所以 i=1ng(x)i=1\sum_{i=1}^n g(x)_i=1

  • MMOE

    yk=hk(i=1ngk(x)ifi(x))y_k=h^k(\sum_{i=1}^n g^k(x)_i f_i(x)), where i=1ngk(x)i=1\sum_{i=1}^n g^k(x)_i=1

    gk(x)=softmax(Wgkx)g^k(x)=softmax(W_{gk}x),其中 xRd,  WgkRn×d, so  gk(x)Rnx \in \mathbb R^d,\ \ W_{gk} \in \mathbb R^{n\times d}, \ so \ \ g^k(x) \in \mathbb R^n

    OMOE 不同之处在于: 每个任务有单独的 Gate network,不共享

Thoughts

  • 是有效的优化,针对任务相关性差的多任务场景,确实能有效涨点
  • 本质是一种参数的堆砌,没看到很创新的点
  • 除了堆参数之外,还有一个很致命的问题,每个 Gate 以及 Expert 都是独立的,实际实现过程中只能使用 for loop 依次计算,效率很低,速度很慢

Experiments

mmoe2.png

mmoe3.png

mmoe4.png