[82] Efficient Language Modeling with Sparse all-MLP (sMLP)

Open dhkim0225 opened this issue 2 years ago • 0 comments

rosinality's comment

mlp with gating으로 lm 학습하기
moe 기반의 sparse 모델, 더 높은 성능을 더 적은 연산량으로 달성
autoregressive lm을 all mlp로 태클
autoregressive이기 때문에 moe routing에 이후 토큰 정보를 쓸 수 없어서 그에 대한 고려가 들어가 있음
spatial mixing에서 어떻게 했는지는 나와있진 않은데 nxn 행렬이니 아마 masking으로 될 것 같네요.

paper

INTRO

all-MLP는 컴퓨팅을 일정하게 유지하면서 모델 용량과 표현력을 크게 증가시킴 저자들은 sparse all-MLP를 제안하여 Transformer 기반 MoE(GShard, Switch Transformer, Base Layers, HASH Layers)와 Dense Transformer 및 all-MLP에 비해 train efficiency가 2배 까지 좋아짐을 보임.

Contributions

spare all-MLP (sMLP) 제안. MoE 를 모든 MLP 와 엮는 건 NLP 에서 처음 있는 시도. 새로운 routing 방식도 제안
좋은 성능

Background

Token-wise Operations

Notations

T == sequence length
H == hidden dim
h == total head number
Y == output
X == input token (T x H)
d == single head 별 dimension (H / h)

token-mixing 을 위해서

transformer 는 self-attention 수행
MLP 모델은 Spatial Gating Unint 사용 (gMLP 에서 채택한 방식) W_s 를 attention score 처럼 생각할 수 있음.

Sparse Expert Models

MoE 가 요즘 흥한다. 여러가지 routing 방법들이 나왔는데, baseline으로 사용하는 방법은 다음과 같다.

N개의 expert.

router weight W_r 을 이용해서, h(x) 를 구한다.

i번째 expert 에 대한 gate-value는 다음과 같이 구한다.

이 중에서 top-k (보통 2개) 만큼 expert 를 고르고, weighted sum 으로 최종 output 을 계산한다.

몇몇 expert 에만 학습이 치중되는 것을 막기 위해,

Switch Transformer나 gshard 는 differential load balancing loss 를 추가해서 device 끼리 communication latency 를 줄였다.
Base Layers 는 balancing loss 대신 balancing assignment algorithm 을 이용해서 좀 더 파이프라인을 단순화 시켰고,
HASH Layers 는 routing weight W_r 을 learning 시키는 대신, random hashing function 을 routing gate 로 사용했다.
논문에는 안 쓰여 있지만 #82 DTS Gate 시도도 있다.

Base Layers 의 fig1. 그냥 token 개수 expert 마다 맞춰주기. Base Layers 사용 방식을 tMoE 라고 부르겠다.

Methods

sMLP

전체적인 pipeline 이다.

N1 개의 dense block 과 N2 개의 sparse block을 갖는다. Base Layers 의 tMoE 로 transformer 의 FFN 대체 self-attn 을 sMoE & Spatial Gating Unit 으로 대체

tMoE 는 token 을 routing 하고, sMoE 는 hidden dimension 을 routing 한다.

Routing in Feature Space

autoregressive lm을 all mlp로 태클해야 한다. autoregressive이기 때문에 MoE routing에 이후 토큰 정보를 쓸 수 없어서 해결책으로 2가지 routing 전략 제안

Deterministic Routing (sMoE)
1. routing strategy 를 배우는 대신, hidden-vector를 일정하게 잘라서 각 expert 에게 보냄.
2. 각각의 expert 를 모두 활용하는 점에서 MHSA 와 닮았음
Partial Prediction (sMoE)
1. 전체 sentence 에서 routing strategy 를 배우는 대신 첫 20% 토큰들로만 routing strategy 를 학습시키고, 나머지 80% 는 prediction으로 나눈다.