Varuna Jayasiri: Attention with Linear Biases (ALiBi)

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

원문: 「Attention with Linear Biases (ALiBi)」, labml.ai 주석 달린 논문 구현, https://nn.labml.ai/transformers/alibi/index.html. 원본 소스는 GitHub의 labml_nn/transformers/alibi/__init__.py이며, 2021년 8월 28일에 처음 올라왔습니다. 이 글은 원문을 한국어로 옮긴 번역문입니다.

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

선형 편향 어텐션ALiBi

이 글은 논문 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation(짧게 학습하고 길게 시험하기: 선형 편향 어텐션으로 입력 길이 외삽하기)에 나온 선형 편향 어텐션ALiBi을 구현한 것입니다.

ALiBi는 위치 인코딩을 쓰는 대신, 어텐션 점수(소프트맥스 앞의 어텐션 로짓)에 편향을 더합니다. 자기회귀 작업에서 시험한 상대 위치 방식으로, 가까운 토큰일수록 편향이 크고 먼 토큰일수록 편향이 작습니다. 편향은 로그 스케일에서 선형으로 줄어들며(소프트맥스 앞이기 때문입니다), 헤드마다 기울기가 다릅니다.

$i$ 번째 토큰의 어텐션 식은 다음과 같습니다.

\begin{aligned} \mathbf{a}_i &= \text{softmax} \bigg( \mathbf{q}_i \mathbf{K}^\top + m \cdot \big[-(i-1), \dots, -1, 0 \big] \bigg) \\ &= \text{softmax} \bigg( \mathbf{q}_i \mathbf{K}^\top + m \cdot \big[0, 1, \dots, (i - 1) \big] \bigg) \end{aligned}

여기서 $\mathbf{q}_i \in \mathbb{R}^d$ 는 $i$ 번째 토큰의 쿼리이고, $K \in \mathbb{R}^{i \times d}$ 는 $i$ 번째까지의 키들이며, $d$ 는 헤드 하나의 특징 수입니다. 위 등식이 성립하는 것은 $\text{softmax}$ 가 평행 이동에 불변이기 때문입니다(모든 원소에 같은 상수를 더해도 결과가 바뀌지 않습니다).

ALiBi 모델의 학습 코드도 있습니다.

import math
from typing import Optional

import torch
from torch import nn

from labml.logger import inspect
from labml_nn.transformers.mha import MultiHeadAttention

헤드마다 다른 기울기 $m$ 구하기

n_heads는 어텐션 층의 헤드 수 $n$ 입니다.

첫 번째 헤드의 기울기는 다음과 같습니다.

\frac{1}{2^{\frac{8}{n}}} = 2^{-\frac{8}{n}}

나머지 헤드들의 기울기는 위와 같은 비율을 공비로 하는 등비수열을 이룹니다.

예를 들어 헤드 수가 $8$ 이면 기울기는 다음과 같습니다.

\frac{1}{2^1}, \frac{1}{2^2}, \dots, \frac{1}{2^8}

def get_slopes(n_heads: int):

n_heads에 가장 가까운 2의 거듭제곱을 구합니다. n_heads가 2의 거듭제곱이 아니면, 먼저 가장 가까운(더 작은) 2의 거듭제곱에 대해 기울기를 계산하고, 그다음 남은 기울기를 덧붙입니다.

    n = 2 ** math.floor(math.log2(n_heads))

2^{-\frac{8}{n}}

    m_0 = 2.0 ** (-8.0 / n)

2^{-1\frac{8}{n}}, 2^{-2 \frac{8}{n}}, 2^{-3 \frac{8}{n}}, \dots

    m = torch.pow(m_0, torch.arange(1, 1 + n))

n_heads가 2의 거듭제곱이 아니면 남은 기울기를 덧붙입니다. 남은 기울기는 $n * 2$ 에 대해 계산하되(앞에서 이미 넣은 기울기는 피합니다), n_heads개까지만 고릅니다.

    if n < n_heads:

2^{-\frac{8}{2n}}

        m_hat_0 = 2.0 ** (-4.0 / n)

$2^{-1\frac{8}{2n}}, 2^{-3 \frac{8}{2n}}, 2^{-5 \frac{8}{2n}}, \dots$ 앞에서 이미 넣은 기울기를 피하려고 $2$ 씩 건너뛴다는 점에 유의하세요.

        m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (n_heads - n), 2))

기울기에 남은 기울기를 이어 붙입니다.

        m = torch.cat([m, m_hat])

    return m

어텐션 편향 행렬 계산하기

n_heads는 어텐션 층의 헤드 수입니다.
mask는 모양이 [seq_len_q, seq_len_k]인 어텐션 마스크입니다.

ALiBi 어텐션 편향을 담은, 모양이 [seq_len_q, seq_len_k, n_heads, ]인 행렬을 돌려줍니다.

@torch.no_grad()
def get_alibi_biases(n_heads: int, mask: torch.Tensor):

헤드마다 기울기 $m$ 을 구합니다.

    m = get_slopes(n_heads).to(mask.device)

거리 $[0, 1, \dots, N]$ 을 계산합니다. 여기서는 마스크를 써서 거리를 계산합니다.

인과 마스크이므로 그냥 $[0, 1, \dots, N]$ 을 써도 됩니다. distance = torch.arange(mask.shape[1], dtype=torch.long, device=mask.device)[None, :]

    distance = mask.cumsum(dim=-1)

둘을 짝지어 곱해 ALiBi 편향 행렬을 얻습니다.

    return distance[:, :, None] * m[None, None, :]

선형 편향 어텐션ALiBi

멀티헤드 어텐션을 재정의합니다.

class AlibiMultiHeadAttention(MultiHeadAttention):

    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
        super().__init__(heads, d_model, dropout_prob)

ALiBi 편향을 캐시해 둘 곳입니다.

        self.alibi_biases = None

query, key, value는 쿼리, 키, 값 벡터의 모음을 담은 텐서입니다. 모양은 [seq_len, batch_size, d_model]입니다.

mask의 모양은 [seq_len, seq_len, batch_size]이고, mask[i, j, b]는 배치 b에서 위치 i의 쿼리가 위치 j의 키와 값에 접근할 수 있는지를 나타냅니다.

    def forward(self, *,
                query: torch.Tensor,
                key: torch.Tensor,
                value: torch.Tensor,
                mask: Optional[torch.Tensor] = None):

ALiBi는 인과 마스크에서만 작동합니다.

        assert mask is not None
        assert mask.shape[0] == mask.shape[1] and mask.shape[2] == 1

query, key, value의 모양은 [seq_len, batch_size, d_model]입니다.

        seq_len, batch_size, _ = query.shape

마스크에 헤드 차원을 더하고 모양을 확인합니다.

        mask = self.prepare_mask(mask, query.shape, key.shape)

어텐션 계산을 위해 query, key, value를 준비합니다. 그러면 모양이 [seq_len, batch_size, heads, d_k]가 됩니다.

        query = self.query(query)
        key = self.key(key)
        value = self.value(value)

어텐션 점수 $Q K^\top$ 을 계산합니다. 모양이 [seq_len, seq_len, batch_size, heads]인 텐서가 나옵니다.

        scores = self.get_scores(query, key)

점수의 크기를 조정합니다. $\frac{Q K^\top}{\sqrt{d_k}}$

        scores *= self.scale

ALiBi 편향이 캐시되어 있지 않으면 만듭니다.

        if self.alibi_biases is None or self.alibi_biases.shape[1] < seq_len:

mask의 모양은 [seq_len, seq_len, 1, 1]입니다.

            self.alibi_biases = get_alibi_biases(scores.shape[-1], mask[:, :, 0, 0])

어텐션 점수에 ALiBi 편향을 더합니다. ALiBi 편향의 모양은 [seq_len, seq_len, n_heads]이고 scores의 모양은 [seq_len, seq_len, batch_size, n_heads]입니다.

        scores += self.alibi_biases[:seq_len, :seq_len, None, :]

마스크를 적용합니다.

        scores = scores.masked_fill(mask == 0, float('-inf'))

키 시퀀스 차원을 따라 $softmax$ 어텐션을 계산합니다. $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$

        attn = self.softmax(scores)

드롭아웃을 적용합니다.

        attn = self.dropout(attn)

값을 곱합니다.

\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V

        x = torch.einsum("ijbh,jbhd->ibhd", attn, value)

여러 헤드를 이어 붙입니다.

        x = x.reshape(seq_len, batch_size, -1)

출력 층입니다.

        return self.output(x)

기울기를 확인해 보는 간단한 테스트 함수입니다.

def _test_alibi():
    inspect(get_slopes(12).tolist(), _n=-1)
    from labml_nn.transformers.utils import subsequent_mask

    mask = subsequent_mask(8)[:, :, 0]
    inspect(mask)

    inspect(get_alibi_biases(12, mask)[:, :, 3], _n=-1)


if __name__ == '__main__':
    _test_alibi()

라이선스 고지

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: 선형 편향 어텐션(ALiBi)Varuna Jayasiri: Attention with Linear Biases (ALiBi)

선형 편향 어텐션ALiBi

헤드마다 다른 기울기 $m$ 구하기

어텐션 편향 행렬 계산하기

선형 편향 어텐션ALiBi

라이선스 고지

선형 편향 어텐션ALiBi

헤드마다 다른 기울기 mmm 구하기

어텐션 편향 행렬 계산하기

선형 편향 어텐션ALiBi

라이선스 고지

헤드마다 다른 기울기 $m$ 구하기