Varuna Jayasiri: Switch Transformer

번역문 · 원저자: Varuna Jayasiri (labml.ai) · 원문: Switch Transformer · 원문 라이선스: MIT (Copyright (c) 2020 Varuna Jayasiri). 원문의 저작권 표시와 MIT 허가 문구는 글 끝 「라이선스 고지」에 그대로 옮겼습니다.

원문: 「Switch Transformer」, labml.ai 주석 달린 논문 구현, https://nn.labml.ai/transformers/switch/index.html. 원본 소스는 GitHub의 labml_nn/transformers/switch/__init__.py이며, 2021년 1월 20일에 처음 올라왔습니다. 이 글은 원문을 한국어로 옮긴 번역문입니다.

라이선스: 이 페이지는 labmlai/annotated_deep_learning_paper_implementations 저장소의 소스 코드 주석으로 만들어지며, 저장소는 MIT 라이선스(Copyright (c) 2020 Varuna Jayasiri)를 따릅니다. 원문 페이지는 왼쪽에 설명, 오른쪽에 코드를 나란히 놓는 형식인데, 여기서는 설명을 먼저 옮기고 바로 아래에 해당 코드를 두었습니다. 코드는 원문 그대로 두고, 코드 옆에 붙은 설명 주석은 본문으로 옮겼습니다.

스위치 트랜스포머

이 글은 논문 Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity(스위치 트랜스포머: 단순하고 효율적인 희소성으로 파라미터 1조 개 모델까지 키우기)를 PyTorch로 작게 구현한 것입니다. 저희 구현은 파라미터가 수백만 개뿐이고, 모델 병렬 분산 학습은 하지 않습니다. GPU 한 장으로 학습하지만, 논문에 설명된 스위칭 개념은 그대로 구현합니다.

스위치 트랜스포머는 토큰에 따라 파라미터들 사이를 바꿔 가며 골라서, 토큰마다 다른 파라미터를 씁니다. 그래서 토큰마다 파라미터의 일부만 선택됩니다. 덕분에 파라미터는 더 많이 가지면서 계산 비용은 더 적게 들일 수 있습니다.

스위칭은 각 트랜스포머 블록의 위치별 피드포워드 네트워크FFN에서 일어납니다. 위치별 피드포워드 네트워크는 완전 연결 층 두 개를 차례로 이어 붙인 것입니다. 스위치 트랜스포머에는 FFN이 여러 개(전문가가 여러 개) 있고, 라우터에 따라 그중 무엇을 쓸지 고릅니다. 라우터의 출력은 각 FFN을 고를 확률들이며, 그중 확률이 가장 높은 것 하나만 골라 계산합니다. 그러니 본질적으로 계산 비용은 FFN이 하나일 때와 같습니다. 저희 구현에서는 모든 것이 GPU 한 장에서 일어나기 때문에, FFN이 많거나 크면 병렬화가 잘 되지 않습니다. 분산 환경이라면 (각각 아주 큰) FFN들을 서로 다른 장치에 하나씩 두게 될 것입니다.

논문은 전문가(FFN)들 사이의 부하를 고르게 맞추기 위한 손실 항을 하나 더 도입하고, 라우팅이 고르지 않을 때 토큰을 버리는 문제도 다룹니다.

Tiny Shakespeare 데이터셋으로 스위치 트랜스포머를 학습하는 학습 코드와 노트북도 있습니다.

Colab에서 열기

import torch
from torch import nn

from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.mha import MultiHeadAttention
from labml_nn.utils import clone_module_list

여러 FFN 사이에서 라우팅하기

class SwitchFeedForward(nn.Module):

capacity_factor는 각 전문가의 용량을, 부하가 이상적으로 고르게 나뉠 때와 견준 배수로 나타낸 것입니다.
drop_tokens는 어떤 전문가에게 용량보다 많은 토큰이 라우팅되었을 때 토큰을 버릴지를 정합니다.
is_scale_prob는 FFN의 입력에 라우팅 확률을 곱할지를 정합니다.
n_experts는 전문가 수입니다.
expert는 전문가 층으로, FFN 모듈입니다.
d_model은 토큰 임베딩의 특징 수입니다.
d_ff는 FFN 은닉층의 특징 수입니다.
dropout은 FFN의 드롭아웃 확률입니다.

    def __init__(self, *,
                 capacity_factor: float,
                 drop_tokens: bool,
                 is_scale_prob: bool,
                 n_experts: int,
                 expert: FeedForward,
                 d_model: int):
        super().__init__()

        self.capacity_factor = capacity_factor
        self.is_scale_prob = is_scale_prob
        self.n_experts = n_experts
        self.drop_tokens = drop_tokens

FFN의 사본들을 만듭니다.

        self.experts = clone_module_list(expert, n_experts)

라우팅 층과 소프트맥스입니다.

        self.switch = nn.Linear(d_model, n_experts)
        self.softmax = nn.Softmax(dim=-1)

x는 스위칭 모듈의 입력으로, 모양은 [seq_len, batch_size, d_model]입니다.

    def forward(self, x: torch.Tensor):

나중에 모양을 되돌리려고 지금 모양을 기록해 둡니다.

        seq_len, batch_size, d_model = x.shape

시퀀스 차원과 배치 차원을 평탄화합니다.

        x = x.view(-1, d_model)

토큰마다 라우팅 확률을 구합니다.

p_i(x) = \frac{e^{h(x)_i}}{\sum^N_j e^{h(x)_j}}

여기서 $N$ 은 전문가 수 n_experts이고, $h(\cdot)$ 는 토큰 임베딩의 선형 변환입니다.

        route_prob = self.softmax(self.switch(x))

가장 큰 라우팅 확률과 그 경로를 구합니다. 확률이 가장 높은 전문가로 라우팅합니다.

        route_prob_max, routes = torch.max(route_prob, dim=-1)

각 전문가로 가는 토큰들의 인덱스를 구합니다.

        indexes_list = [torch.eq(routes, i).nonzero(as_tuple=True)[0] for i in range(self.n_experts)]

출력을 담을 빈 텐서를 초기화합니다.

        final_output = x.new_zeros(x.shape)

각 전문가의 용량입니다.

\mathrm{expert\;capacity} = \frac{\mathrm{tokens\;per\;batch}}{\mathrm{number\;of\;experts}} \times \mathrm{capacity\;factor}

        capacity = int(self.capacity_factor * len(x) / self.n_experts)

각 전문가에게 라우팅된 토큰 수입니다.

        counts = x.new_tensor([len(indexes_list[i]) for i in range(self.n_experts)])

버린 토큰을 담을 빈 리스트를 초기화합니다.

        dropped = []

drop_tokens가 True일 때만 토큰을 버립니다.

        if self.drop_tokens:

전문가마다 토큰을 버립니다.

            for i in range(self.n_experts):

전문가가 용량을 넘지 않았으면 건너뜁니다.

                if len(indexes_list[i]) <= capacity:
                    continue

버리기 전에 인덱스를 섞습니다.

                indexes_list[i] = indexes_list[i][torch.randperm(len(indexes_list[i]))]

용량을 넘는 토큰들을 버린 토큰으로 모읍니다.

                dropped.append(indexes_list[i][capacity:])

전문가의 용량까지만 토큰을 남깁니다.

                indexes_list[i] = indexes_list[i][:capacity]

전문가 FFN들의 출력을 구합니다.

        expert_output = [self.experts[i](x[indexes_list[i], :]) for i in range(self.n_experts)]

최종 출력에 채워 넣습니다.

        for i in range(self.n_experts):
            final_output[indexes_list[i], :] = expert_output[i]

버린 토큰은 그대로 통과시킵니다.

        if dropped:
            dropped = torch.cat(dropped)
            final_output[dropped, :] = x[dropped, :]

        if self.is_scale_prob:

전문가 출력에 확률을 곱합니다. $y = p_i(x) E_i(x)$

            final_output = final_output * route_prob_max.view(-1, 1)
        else:

값의 크기는 바꾸지 않되, 그래디언트가 흐르도록 $\frac{p}{\hat{p}} = 1$ 을 곱합니다(저희가 실험해 본 방식입니다).

            final_output = final_output * (route_prob_max / route_prob_max.detach()).view(-1, 1)

최종 출력의 모양을 [seq_len, batch_size, d_model]로 되돌립니다.

        final_output = final_output.view(seq_len, batch_size, d_model)

다음을 돌려줍니다.

최종 출력
각 전문가에게 라우팅된 토큰 수
각 전문가의 확률 합
버린 토큰 수
선택된 전문가들의 라우팅 확률

이 값들은 부하 균형 손실과 기록에 쓰입니다.

        return final_output, counts, route_prob.sum(0), len(dropped), route_prob_max

스위치 트랜스포머 블록

일반 트랜스포머 블록과 같되, 스위치 피드포워드 모듈이 내놓는 추가 출력을 처리합니다.

class SwitchTransformerLayer(nn.Module):

d_model은 토큰 임베딩 크기입니다.
attn은 어텐션 모듈입니다.
feed_forward는 피드포워드 모듈입니다(여기서는 스위칭 모듈입니다).
dropout_prob는 셀프 어텐션과 FFN 뒤에 드롭아웃할 확률입니다.

    def __init__(self, *,
                 d_model: int,
                 attn: MultiHeadAttention,
                 feed_forward: SwitchFeedForward,
                 dropout_prob: float):
        super().__init__()
        self.size = d_model
        self.attn = attn
        self.feed_forward = feed_forward
        self.dropout = nn.Dropout(dropout_prob)
        self.norm_self_attn = nn.LayerNorm([d_model])
        self.norm_ff = nn.LayerNorm([d_model])

    def forward(self, *,
                x: torch.Tensor,
                mask: torch.Tensor):

셀프 어텐션을 하기 전에 벡터를 정규화합니다.

        z = self.norm_self_attn(x)

셀프 어텐션을 거칩니다. 즉, 키와 값이 자기 자신에게서 나옵니다.

        self_attn = self.attn(query=z, key=z, value=z, mask=mask)

셀프 어텐션 결과를 더합니다.

        x = x + self.dropout(self_attn)

피드포워드를 위해 정규화합니다.

        z = self.norm_ff(x)

스위칭 피드포워드 네트워크를 거칩니다.

        ff, counts, route_prob, n_dropped, route_prob_max = self.feed_forward(z)

피드포워드 결과를 다시 더합니다.

        x = x + self.dropout(ff)

        return x, counts, route_prob, n_dropped, route_prob_max

스위치 트랜스포머

class SwitchTransformer(nn.Module):

    def __init__(self, layer: SwitchTransformerLayer, n_layers: int):
        super().__init__()

트랜스포머 층의 사본들을 만듭니다.

        self.layers = clone_module_list(layer, n_layers)

마지막 정규화 층입니다.

        self.norm = nn.LayerNorm([layer.size])

    def forward(self, x: torch.Tensor, mask: torch.Tensor):

트랜스포머 층을 하나씩 거칩니다.

        counts, route_prob, n_dropped, route_prob_max = [], [], [], []
        for layer in self.layers:
            x, f, p, n_d, p_max = layer(x=x, mask=mask)
            counts.append(f)
            route_prob.append(p)
            n_dropped.append(n_d)
            route_prob_max.append(p_max)

마지막으로 벡터를 정규화합니다.

        x = self.norm(x)

        return x, torch.stack(counts), torch.stack(route_prob), n_dropped, torch.stack(route_prob_max)

라이선스 고지

Copyright (c) 2020 Varuna Jayasiri

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Delete your account?계정을 지울까요?

Varuna Jayasiri: 스위치 트랜스포머Varuna Jayasiri: Switch Transformer

스위치 트랜스포머

여러 FFN 사이에서 라우팅하기

스위치 트랜스포머 블록

스위치 트랜스포머

라이선스 고지