🇧🇷 Explicando Transformers

Em 2017, a Google Brain lançou o artigo “Attention is All you Need”, que introduziu para o mundo o Transformer, uma arquitetura para redes neurais para transdução de sequências baseada em atenção que permitiu a criação de modelos que superaram todos os outros modelos anteriores em tradução entre idiomas.

Anos depois, essa arquitetura deu origem a muitas alternativas que atingiram o Estado da Arte múltiplas em outras aplicações, em especial na criação de modelos de linguagem, que popularizaram fortemente o uso de IA Generativa para muitas aplicações de processamento de linguagem natural. No momento de escrita desse post, são variações do Transformer como os Generative Pretrained Transformers (GPTs) que estão sendo usadas por trás de inteligências artificiais avançadas, como ChatGPT, Claude e Gemini, mostrando o potencial dessa arquitetura e o seu legado na história da inteligência artificial.

Para ajudar a entender o que os Transformers têm de especial, esse post vai explicar como funcionam os Transformers, através de uma implementação de exemplo construída do zero em Python usando PyTorch, e assumindo poucos pré-requisitos.

Embora já existam muitas implementações de Transformers disponíveis na internet, durante o momento de escrita desse post, não existiam muitos artigos explicando o tópico em Português de forma didática, o que motivou a escrita.

Infelizmente, se absolutamente nenhum pré-requisito for assumido, o conteúdo ficará extenso demais para ser feito de uma vez (embora talvez seja possível escrever outros posts com os pré-requisitos no futuro). Por isso, para entender tudo que será explicado, é importante entender os conteúdos a seguir:

• Como funcionam operações matriciais como produto interno, produto matricial, transposição de matrizes, etc.
• Como funcionam os componentes de redes neurais feed-forward
• Como programar em Python ao nível de criar classes e objetos e interagir com bibliotecas terceiras
• Como usar os componentes básicos do PyTorch, como tensores, dispositivos, módulos, autograd, otimizadores, etc.
• Porque realizar operações em lote usando operações matriciais podem ser muito mais rápido que realizá-las individualmente.

Além disso, a seção de recursos adicionais possui alguns conteúdos em Inglês que explicam esses mencionados.

Objetivo

De forma geral, Transformers podem ser usados para modelar aplicações de transdução de sequências, que são aplicações onde é necessário gerar uma sequência de elementos a partir de outra. Isso faz com que muitas aplicações possam ser modeladas como uma transdução de sequências, como tradução automática, geração textual, sumarização de textos e síntese de moléculas.

Usando uma definição mais formal, um modelo de transdução de sequências relaciona uma sequência ordenada $s = <s_1,s_2,\cdots,s_{n}>$, composta de elementos de um conjunto enumerável $S$ a outra sequência $t = <t_1,t_2,\cdots,t_{m}>$, composta de elementos de um conjunto enumerável $T$.

Para simplificar a explicação sobre o funcionamento dos Transformers, muitas explicações usarão exemplos pensando em uma aplicação específica. Dada a importância dos Transformers para o progresso de aplicações como os Large Language Models, os exemplos desse post serão baseados no contexto da criação de um modelo de linguagem.

O objetivo de um modelo de linguagem é determinar os próximos caracteres de um texto a partir dos anteriores, assim como fazem os corretores automáticos usados nos teclados digitais de smartphones, por exemplo.

Adaptando a definição anterior, um modelo de linguagem relaciona um texto $s = <s_1,s_2,\cdots,s_{n}>$, composto de caracteres do vocabulário $S$ a outro texto $s' = <s'_1,s'_2,\cdots,s'_{m}>$, também composto de caracteres do vocabulário $S$. Note que $S$ e $T$ serem iguais é uma especificidade de aplicações como modelos de linguagem. Em tradução automática, por exemplo, os textos podem estar em idiomas diferentes, logo, os vocabulários também podem ser diferentes.

Observação

Além de criar um modelo de linguagem de forma geral, um grande desafio no desenvolvimento de redes neurais é maximizar a eficiência computacional desses modelos. Embora seja possível implementar um Transformer usando apenas estruturas como variáveis, listas e laços de repetição, esse tipo de implementação pode gerar modelos lentos demais para treinar e usar na prática pelo grande número de computações, que de forma ingênua, são realizadas de forma ineficiente. Portanto, é necessário levar em consideração o uso de técnicas de programação paralela desde o início da implementação do modelo.

Em geral, a forma mais prática de implementar paralelismo em redes neurais é aproveitar a capacidade dos dispositivos modernos de realizar operações matemáticas sobre vetores, matrizes e tensores em paralelo de forma muito eficiente. Portanto, mudando a representação dos dados durante a execução das operações, é possível gerar algoritmos eficientes de forma simples, embora isso exija que os algoritmos sejam definidos levando em consideração dados multidimensionais desde a sua concepção.

Um pouco de história

Criar modelos eficientes para transdução de sequências foi um problema em aberto por muitos anos para a comunidade científica, onde o desafio era resolver problemas comuns ao treino de arquiteturas alternativas aos Transformers. Mesmo com a criação de técnicas avançadas como os Transformers, modelos satisfatórios desse tipo ainda não existem para diversas aplicações.

Técnicas baseadas em redes neurais recorrentes possuiam a melhor performance em aplicações como tradução textual, como foi o caso da arquitetura proposta pela Google para o Google Tradutor em 2014 no influente artigo “Sequence to Sequence Learning with Neural Networks”). Embora explicar esse tipo de arquitetura esteja fora do escopo desse post, existiam problemas comuns ao processo de treinamento de redes neurais recorrentes que encorajaram o desenvolvimento de alternativas:

1. A sua natureza recursiva pode causar problemas de gradient vanishing, efeito que ocorre quando os valores dos gradientes gerados via backpropagation são pequenos demais para que o modelo consiga convergir para o mínimo global até o final do treinamento, praticamente parando antes do valor da função de perda se tornar próximo do mínimo global.

2. A ordem sequencial das operações envolvidas nesse tipo de modelo pode tornar a inferência muito lentas, o que pode tornar o treinamento e uso posterior do modelo inviável.

3. A incapacidade do modelo de entender o significado de uma palavra em uma frase com base no seu contexto pode comprometer o texto gerado.

Esses problemas eram críticos para o uso desses modelos em muitas aplicações, o que incentivou a pesquisa sobre a criação de arquiteturas alternativos que convergissem mais rapidamente e em modelos com melhor performance. Em especial, a criação de modelos capazes de resignificar palavras com base no seu contexto para resolver parte desses problemas levou ao estudo de mecanismos de atenção, que têm justamente esse objetivo.

Assim como grande parte da pesquisa em redes neurais, muitas das decisões e convenções que serão usadas são baseadas no que levou aos melhores resultados nos experimentos conduzidos. Porém, no caso dos Transformers, também existem escolhas tomadas com o objetivo de garantir que os algoritmos são eficientes computacionalmente em treino e inferência e não sofram dos mesmos problemas de convergência que as alternativas.

Nesse sentido, o grande diferencial dos Transformers é o fato de serem baseados apenas em redes neurais feed-forward e mecanismos de atenção, o que leva a modelos capazes de resolver os problemas mencionados anteriormente em diversas aplicações e atingir a melhor performance até então em transdução de sequências no processo.

Essa conclusão explica o nome do artigo: do ponto de vista da arquitetura do modelo, redes neurais recorrentes não são necessárias para criar modelos eficientes, apenas mecanismos de atenção, ou seja, “Attention is All you Need”.

Representação dos dados

Para conseguirmos trabalhar com textos, é necessário definir uma representação numérica equivalente a um texto para que essa representação possa ser usada pelos Transformers. A abordagem que a maioria das técnicas usa para realizar esse processo é converter esses valores em tokens e embeddings.

Tokens

Como textos são representados por um alfabeto finito e conhecido, é possível enumerar todos os caracteres desse alfabeto e criar uma função que os associa a uma representação numérica única. Essas representações numéricas são conhecidas como tokens, e a função como tokenizer.

Além de enumerar caracteres individuais, é possível enumerar sequências de caracteres, formando palavras ou n-gramas (sequências ordenadas de n caracteres), gerando tokenizers com ainda mais tokens. Essa é uma prática realizada muito frequentemente em modelos de linguagem, com o objetivo de melhorar a performance desses modelos de forma geral. usando algortimos como Byte Pair Encoding, WordPiece e SentencePiece, é possível criar tokenizers com centenas de milhares de tokens.

Porém, como o modelo de linguagem assume que o texto já passou pelo tokenizer previamente, a escolha do método de tokenização não afeta a arquitetura de nenhuma forma. Por isso, para simplificar a explicação, usaremos um tokenizer simples composto apenas pelos caracteres imprimíveis da tabela ASCII, disponível no objeto printable do módulo string em Python.

Uma sequência de tokens gerada a partir de um texto pode ser representanda usando um vetor com os valores de cada token em ordem. Concatenando esses vetores como linhas, é possível representar um batch de textos usando uma matriz, que será a dimensão esperada pelo modelo para maximizar a sua eficiência computacional.

Além disso, o tokenizer adiciona dois tokens especiais à sequência. O primeiro token, <bos> (beginning of sentence), representa o início do texto, e o segundo token, <eos> (end of sentence), representa o final do texto.

Nos exemplos, <bos> e <eos> sempre serão representados numericamente por 1 e 2, respectivamente, logo, as representações numéricas mencionadas anteriormente começam a partir de 3.

Os tokens <bos> e <eos> também podem ser usados para separar a sequência recebida da sequência gerada, o que será útil no treinamento dos Transformers futuramente. Por exemplo, se o modelo recebeu o texto “cachorro” e gerou o texto “dog”, essas sequências pode ser descritas como sendo uma sequência só:

from string import printable

string_to_int = {
    '<bos>': 1,
    '<eos>': 2,
}

string_to_int.update({
    char: index + 2
    for index, char in enumerate(printable, 1)
})

special_tokens = {'<bos>', '<eos>'}

def split(text: str) -> Generator[str, None, None]:
    index = 0

    while index < len(text):
        token = text[index]
        for special_token in special_tokens:
            if text.startswith(special_token, index):
                token = special_token
                break

        yield token

        index += len(token)

def tokenize(self: Self, text: str) -> torch.Tensor:
    tokens = [string_to_int[token] for token in split(text)]
    tokens = torch.tensor(tokens)
    return tokens

Padding

Como introduzido anteriormente, uma forma de desenvolver modelos rápidos de forma simples é maximizar o número de operações realizadas em batch usando operações tensoriais.

Um texto após ser transformado em sequência de tokens pode ser representado como um vetor, e ao realizar essa transformação em todos os textos de um lote e concatená-los, é possível criar uma matriz, que representa um lote de textos e permite a aceleração de operações tensoriais nas próximas etapas. Porém, concatenar esses vetores para criar uma matriz assume que todos os elementos do lote possuem o mesmo comprimento, mas como os textos originais podem ter comprimentos diferentes, essa premissa não é sempre válida.

Uma forma de garantir que todos os vetores de um batch sempre terão o mesmo comprimento é realizar padding nos vetores, processo onde tokens especiais são adicionados repetidamente a cada sequência até que todas tenham $n$ elementos.

Esse token especial é denominado padding token, representado por <pad> e nos exemplos, sempre será representado numericamente por 0. Também assuma que todos um batch sempre tem $b$ elementos.

Por fim, existem algumas formas de determinar $t$, o número de elementos que cada sequência terá após o padding:

1. Definir $t$ como sendo o comprimento da maior sequência do batch.
2. Definir $t$ como sendo um valor constante e arbitrário.

Por exemplo, um lote de textos transformado em tokens e usando padding seguindo a opção 1 seria feito da seguinte forma:

De forma geral, o aumento de velocidade obtido pelo padding é grande o suficiente para justificar o uso adicional de memória, já que o tempo gasto é frequentemente um limitante maior que a memória gasta por tipo de rede neural durante o treinamento e inferência.

As duas alternativas podem impactar a velocidade das operações dependendo do tipo de hardware usado. Em CPUs e GPUs, a opção 1 pode ser mais vantajosa por economizar memória, já que nesse caso a opção 2 não gera nenhum aumento de velocidade. Já em TPUs, operações sobre batches de tamanhos fixos podem ser mais rápidas que as mesmas operações em batches de tamanhos diferentes, logo a opção 2 pode ser mais vantajosa.

char_ids = {
    '<pad>': 0,
    '<bos>': 1,
    '<eos>': 2,
}

char_ids.update({
    char: index + 3
    for index, char in enumerate(printable)
})

special_tokens = {'<pad>', '<bos>', '<eos>'}

def pad(
    self: Self,
    tokens: torch.Tensor,
    amount: int,
    fill_value: int,
    side: Literal["left", "right"],
) -> torch.Tensor:
    if amount == 0:
        return tokens

    padding = torch.full(
        size=(amount,),
        fill_value=fill_value,
    )
    padded_tensor = (tokens, padding) if side == "right" else (padding, tokens)
    padded_tensor = torch.cat(padded_tensor)

    return padded_tensor

def batch_encode(
    texts: list[str],
    side: Literal["left", "right"] = "right",
    strategy: Literal["max", "fixed"] = "max",
    amount: int | None = None,
    truncate: bool = False,
) -> torch.Tensor:
    token_lists = [encode(text) for text in texts]
    lengths = [len(tokens) for tokens in token_lists]

    max_length = max(lengths)
    max_length = max_length if amount is None else max(max_length, amount)

    token_lists = [
        pad(
            tokens=tokens,
            amount=max_length - length,
            fill_value=0,
            side=side,
        )
        for tokens, length in zip(token_lists, lengths)
    ]
    token_lists = torch.stack(token_lists)

    if strategy == "fixed" and truncate:
        token_lists = token_lists[:, :amount]

    return token_lists

Por curiosidade, o módulo torch.nested permite representar matrizes a partir de uma lista de vetores de comprimentos variados. Porém, esse módulo ainda está em fase experimental, e usar esse recurso pode deixar as operações mais lentas que as outras alternativas.

Truncação

Mesmo ao escolher a opção 1, também é comum definir um tamanho máximo que as sequências podem ter, e caso uma sequência tenha originalmente mais tokens que esse limite, ela será truncada, fazendo com seus últimos elementos sejam descartados.

Em ambos os casos, o valor de $t$ geralmente é escolhido com base na memória disponível ou determinando empiricamente o valor para o comprimento de uma sequência onde, em média, os modelos sendo treinados não conseguem considerar toda a sequência recebida durante a geração.

Implementação completa do Tokenizer

class Tokenizer:
    def __init__(
        self: Self,
        vocab: set[str],
        config: TokenizerConfig,
    ) -> None:
        super().__init__()
        self.vocab = vocab
        self.config = config

        self.string_to_int = {
            self.config.pad_token: 0,
            self.config.bos_token: 1,
            self.config.eos_token: 2,
        }
        self.string_to_int.update(
            {char: (index + 3) for index, char in enumerate(self.vocab)}
        )

        self.pad_token_int = self.string_to_int[self.config.pad_token]
        self.bos_token_int = self.string_to_int[self.config.bos_token]
        self.eos_token_int = self.string_to_int[self.config.eos_token]

        self.vocab_size = len(self.string_to_int)

        self.int_to_string = {
            0: self.config.pad_token,
            1: self.config.bos_token,
            2: self.config.eos_token,
        }
        self.int_to_string.update(
            {(index + 3): char for index, char in enumerate(self.vocab)}
        )

    def split(self: Self, chars: str) -> Generator[str, None, None]:
        index = 0
        while index < len(chars):
            token = chars[index]
            for special_token in self.config.special_tokens:
                if chars.startswith(special_token, index):
                    token = special_token
                    break

            yield token
            index += len(token)

    def encode(self: Self, text: str) -> torch.Tensor:
        tokens = [self.string_to_int[token_string] for token_string in self.split(text)]
        tokens = torch.tensor(tokens)
        return tokens

    def pad(
        self: Self,
        tokens: torch.Tensor,
        amount: int,
        fill_value: int,
        side: Literal["left", "right"],
    ) -> torch.Tensor:
        if amount == 0:
            return tokens

        padding = torch.full(
            size=(amount,),
            fill_value=fill_value,
        )
        padded_tensor = (tokens, padding) if side == "right" else (padding, tokens)
        padded_tensor = torch.cat(padded_tensor)

        return padded_tensor

    def batch_encode(
        self: Self,
        texts: list[str],
        side: Literal["left", "right"] = "right",
        strategy: Literal["max", "fixed"] = "max",
        amount: int | None = None,
        truncate: bool = False,
    ) -> torch.Tensor:
        token_lists = [self.encode(text) for text in texts]
        lengths = [len(tokens) for tokens in token_lists]

        max_length = max(lengths)
        max_length = max_length if amount is None else max(max_length, amount)

        token_lists = [
            self.pad(
                tokens=tokens,
                amount=max_length - length,
                fill_value=self.pad_token_int,
                side=side,
            )
            for tokens, length in zip(token_lists, lengths)
        ]
        token_lists = torch.stack(token_lists)

        if strategy == "fixed" and truncate:
            token_lists = token_lists[:, :amount]

        return token_lists

    def decode(self: Self, tokens: torch.Tensor) -> str:
        outputs = [self.int_to_string[token] for token in tokens.tolist()]
        outputs = "".join(outputs)
        return outputs

    def batch_decode(self: Self, tokens: torch.Tensor) -> list[str]:
        outputs = [self.decode(item) for item in tokens]
        return outputs

    def batch_shift(
        self: Self,
        input_tokens: torch.Tensor,
        output_tokens: torch.Tensor,
    ) -> str:
        outputs = torch.cat((input_tokens[:, 1:], output_tokens), dim=1)
        return outputs

    def add_special_tokens(self: Self, text: str) -> str:
        return self.config.bos_token + text + self.config.eos_token

Embeddings

A representação numérica dos tokens é uma forma simples de converter caracteres para valores numéricos. Porém, usá-la diretamente como espaço de representação dos elementos da sequência recebida durante o treinamento de modelos pode criar vieses indesejados durante o treinamento.

Esse viés ocorre porque usar a representação numérica literal fará com que o modelo trate elementos com o token $n$ como sendo valores menores que elementos com o token $n+1$, já que ambos são inteiros. Embora essa relação não seja verdadeira, é possível que o modelo encontre padrões que não existem nos conjunto de dados de treinamento, dificultando o processo. Ao transformar esses valores em vetores de dimensão $t$, se o espaço for adequado, essa relação de ordem não estará presente.

Embeddings são representações do espaço de uma variável (por exemplo, os inteiros para a representação numérica dos tokens) em outro espaço. Em aprendizado de máquina, esse espaço geralmente possui $d$ dimensões, onde $d$ é um hiperparâmetro diferente das dimensões do espaço original e tanto $d$ quanto o espaço são escolhidos com o objetivo de representar os valores originais de outra forma que auxilie o treinamento de modelos.

Independentemente da técnica usada para determinar os embeddings para cada elemento, o batch de sequências recebidas será transformado de uma matriz de dimensão $b \times t$ para cada um tensor de dimensão $b \times t \times d$.

Usando uma das formas mais simples de se transformar tokens em embeddings de forma eficiente, primeiro é necessário aplicar one-hot encoding sobre cada elemento, transformando-os vetores esparsos (que possui grande dimensionalidade mas muitos valores nulos).

Se $|T|$ é o tamanho do vocabulário, a representação do texto “baba” usando One-Hot Encoding será:

One-Hot Encoding é usado por permitir acessar elementos individuais de uma matriz a partir de um batch de índices, que no caso, será a sequência recebida. Isso ocorre porque ao multiplicar uma das linhas da sequência após One-Hot Encoding, de dimensões $1 \times |T|$, por uma matriz de dimensões $|T| \times d$, o resultado será a $m$-ésima linha dessa matriz. Essa matriz é denominada matriz de consulta, já que esse comportamento é similar a acessar uma tabela de consulta, ou acessar os elementos de uma lista ou array através de índices.

Logo, ao multiplicar a sequência inteira após One-Hot Encoding, de dimensões $t \times d$, pela matriz de consulta, o resultado será as suas linhas concatenadas de acordo com a ordem dos tokens da sequência recebida.

No exemplo, vamos considerar que o embedding de cada token possui sempre o mesmo valor após essa transformação. Nesse caso, é possível concatenar todos os $|T|$ embeddings possíveis em ordem para criar uma matriz de consulta, e a combinação dessas operações pode representar o processo de transformação da sequência em embeddings.

Além de ser acelerável usando álgebra linear, essa forma de acesso aos embeddings permite que os valores dos embeddings possam ser otimizados durante o treinamento do modelo para melhorar a sua performance. Para isso, basta tratar a matriz de consulta como se fossem os pesos de uma camada linear em uma rede neural feed-forward, e otimizar o seu valor a partir da função de perda do modelo através de backpropagation. Assim, é possível encontrar um espaço ótimo para representar os tokens para o conjunto de dados de treinamento, sendo necessário definir apenas o valor de $d$.

Além da performance, os vetores encontrados nesse espaço são densos, já que $d$, que costuma ser um valor na casa dos milhares, geralmente é muito menor do que $|T|$, que pode chegar às centenas de milhares de tokens no contexto de Large Language Models. Isso faz com que o modelo requira muito menos memória durante o seu uso e evita problemas causados pela dimensionalidade elevada presente ao treinar modelos usando One-Hot Encoding diretamente.

Em PyTorch, esse componente está encapsulado na classe torch.nn.Embedding.

embed_dim = 512

embedder = nn.Embedding(
    num_embeddings=len(printable) + 1,
    embedding_dim=embed_dim,
)

embedding = embedder(42)

Atenção

No contexto de redes neurais para transdução de sequências, o uso de atenção, de forma geral, se refere à capacidade do modelo de entender o contexto de cada elemento da sequência ao gerar um novo valor para cada elemento recebido.

Nesse processo, os pesos de atenção são gerados, uma sequência intermediária representando a intensidade com que cada elemento recebido deve ser utilizado para gerar os novos elementos individualmente. Ao combinar os pesos de atenção com os elementos originais, a sequência gerada é obtida.

O uso de mecanismos de atenção dessa forma permite que modelos compreendam o contexto onde cada palavra está inserida usando menos operações, o que acelera treinamentos e pode otimizar os modelos (por evitar que não entendam textos ambíguos sem considerar o contexto).

Por exemplo, considere o texto a seguir:

João pensou: O trem estava cheio, mas ele conseguiu uma cadeira livre.

Por simplicidade, ignore os tokens especiais e suponha que cada palavra é representada por um embedding (pelo processo abstraído por $\text{Emb}(X)$), transformando o texto na sequência a seguir:

Então, um mecanismo de atenção recebe essa sequência e gera outra de mesmo comprimento, onde o embedding que representa a palavra “ele” será composto de embeddings mais próximos do trecho que compõe a palavra “João” do que da palavra “ônibus”:

Self-Attention

Todos os mecanismos que serão usados na arquitetura dos Transformers são baseados apenas em combinar os elementos da sequência recebida entre si para determinar a atenção para cada elemento.

Porém, nem todos os mecanismos funcionam dessa forma. Outros dependem de recursos externos para determinar a atenção, como outras variáveis ou modelos. Por isso, é dito que os mecanismos explicados usam Self-Attention, o que significa que assumem que a atenção para cada elemento pode ser determinada corretamente sem a necessidade de usar outros recursos externos.

ELI5: Queries, Keys e Values

Para explicar conceitualmente o funcionamento desses mecanismos, embora os termos dessas operações sejam numericamente iguais no primeiro momento, alguns termos receberão nomes abstratos diferentes.

De forma reducionista, um mecanismo de atenção funciona de como um dicionário em Python, onde chaves (keys ou $K$) são associadas a valores (values ou $V$) e é possível recuperar um valor posteriormente a partir da sua chave, denominada consulta (query ou $Q$). A analogia é que no contexto do mecanismo, quem exerce esses papéis são:

• Query: Um elemento da sequência recebida, como a palavra “ele”.
• Keys: Serão todos os elementos originais da sequência recebida
• Value: Será algum elemento da sequência recebida.
  • Caso seja necessário entender algum contexto, será um elemento que melhor representa a query, como a palavra “João” no caso da palavra “ele”.
  • Caso contrário, será a própria query.

O exemplo a seguir ilustra essa analogia:

sequence = keys = values = [
    'João',
    'pensou',
    'O',
    'trem',
    'estava',
    'cheio',
    'mas',
    'ele',
    ...
]

mechanism = AttentionMechanism(keys, values)
query = 'ele'

assert mechanism[query] == 'João'

Uma das grandes diferenças entre as duas estruturas é que retornar apenas um elemento para cada query pode não ser o suficiente para representar o contexto corretamente em situações mais ambíguas. Por exemplo:

Vi João e Maria ontem. Eles estavam juntos.

Nesse caso, a palavra “Eles” não pode ser compreendida corretamente usando apenas um value. Por isso, um mecanismo de atenção não retornará apenas um valor, e sim uma fração do quanto cada um dos values possíveis deve ser considerado:

sequence = keys = values = [
    'Vi',
    'João',
    'e',
    'Maria',
    'ontem.',
    'Eles',
    ...
]

mechanism = AttentionMechanism(keys, values)
query = 'ele'

assert mechanism[query] == {
    'Vi': 0.01,
    'João': 0.45,
    'e': 0.01,
    'Maria': 0.45,
    'ontem.': 0.01,
    'Eles': 0.01,
    ...
}

Essas frações serão os pesos de atenção e são uma sequência normalizada, ou seja, a soma de todos os elementos é 1. Como os pesos são normalizados e nos Transformers as sequências recebidas serão representadas por embeddings, é possível aplicar uma média ponderada usando os pesos para gerar um novo elemento, que será a sequência gerada pelo mecanismo de atenção.

Os pesos precisam ser normalizados para que a atenção que pode ser distribuída seja finita e o mecanismo funcione corretamente. Então, Se o score de um elemento for relativamente maior, o score de pelo menos um dos outros elementos será relativamnnte menor de forma proporcional para manter essa propriedade.

Como nos Transformers será necessário transformar cada elemento da sequência recebida, o mecanismo também poderá ser acelerado ao ser realizado em batch, fazendo com que as queries também sejam iguais à sequência recebida.

Por isso, embora queries, keys e values partam da mesma sequência inicialmente, é importante separar o seu papel em cada parte do mecanismo.

Mecanismos de atenção

Dot-Product Attention (DPA)

Nesse mecanismo, os pesos de atenção para cada elemento são determinados a partir do produto interno entre queries e keys.

Calculando a nova sequência em batch, o DPA poderá ser definido como:

Porém, o produto interno de dois vetores não está contido entre 0 e 1, e sim entre $-\infty$ e $\infty$. Dessa forma, o mecanismo poderá atribuir atenção infinitamente entre todos os elementos, o que pode enviesar o modelo e inviabilizar o uso do mecanismo. Para corrigir isso e normalizar a atenção, é aplicada a função softmax sobre os pesos.

Scaled Dot-Product Attention (SDPA)

Essa variação do DPA adiciona um termo de normalização sobre os pesos para estabilizar os gradientes gerados pelo mecanismo durante a etapa de backpropagation:

A origem da normalização por $\sqrt{d}$ é empírica. Antes dos Transformers, já se observava experimentalmente que normalizar a variância do gradiente gerado por camadas ocultas evita problemas de gradient vanishing e neurônios mortos durante o treinamento de modelos, característica que foi adaptada para a arquitetura posteriormente.

def apply_scaled_dot_product_attention(
    queries: torch.Tensor,
    keys: torch.Tensor,
    values: torch.Tensor,
) -> torch.Tensor:
    keys = keys.transpose(2, 3)
    scores = queries @ keys / (split_embed_dim**0.5)

    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))

    weights = F.softmax(scores, dim=3)
    outputs = weights @ values

    return outputs

Projeções lineares

Uma das formas com que os Transformers tornam o SDPA mais eficiente é projetar linearmente as queries, keys e values para espaços diferentes, multiplicando-as por matrizes de parâmetros treináveis (denotadas como $W^Q$, $W^K$ e $W^V$). Isso faz com que os valores das queries, keys e values se tornem diferentes entre si, e durante o treino, esses parâmetros sejam otimizados para otimizar a forma como o transformer converte os embeddings originais em outro espaço com as mesmas dimensões, mas que representam melhor o valor de cada token no contexto onde estão inseridos.

queries_projection = nn.Linear(embed_dim, embed_dim, bias=False)
keys_projection = nn.Linear(embed_dim, embed_dim, bias=False)
values_projection = nn.Linear(embed_dim, embed_dim, bias=False)

queries = queries_projection(queries)
keys = keys_projection(keys)
values = values_projection(values)

embeddings = apply_scaled_dot_product_attention(queries, keys, values)

Multihead Attention (MHA)

Essa variação do SDPA consiste em aplicar o mecanismo $h$ vezes sobre os embeddings para combinar os resultados dessas aplicações, onde o número de cabeças $h$ é um hiperparâmetro do modelo.

Como os pesos no SDPA são normalizados, não é possível prestar atenção em todos os elementos simultaneamente. O objetivo do uso de MHA é fazer com que cada SDPA foque em um tipo de padrão diferente no representar do contexto para criar um modelo melhor durante o treinamento.

Usando uma analogia, MHA seria o equivalente a ler um texto $h$ vezes, permitindo prestar atenção em partes diferentes do texto a cada vez para entender melhor o contexto de cada palavra. No mecanismo, a diferença é que todas as leituras podem ser feitas ao mesmo.

Embora esse mecanismo seja eficaz em otimizar modelos, se for implementado da forma literal, é necessário calcular SDPA $h$ vezes, tornando o MHA $h$ vezes mais lento, algo que não é desejado, já que um dos dos objetivos dos Transformers é criar modelos eficientes computacionalmente. Por isso, é usado o algoritmo alternativo a seguir para fazer com que a escalabilidade do mecanismo não seja afetada pelo número de cabeças:

1. Dividir as queries, keys e values em $h$ trechos contíguos, transformando as dimensões do batch de $b \times t \times d$ para $b \times t \times h \times \frac{h}{d}$.
2. Reordenar a ordem dos elementos das queries, keys e values, transformando as dimensões do batch de $b \times t \times h \times \frac{h}{d}$ para $b \times h \times t \times \frac{h}{d}$.
3. Aplicar SDPA no batch $h$ vezes, usando projeção diferentes para cada cabeça.
4. Concatenar os embeddings de cada cabeça.
5. Restaurar a ordem dos elementos no batch gerado, restaurando as dimensões para $b \times t \times d$.
6. Aplicar uma projeção linear $W^O$ nos embeddings concatenados.

Após a etapa 3, o algoritmo pode ser descrito pela equação a seguir:

Dessa forma, SDPA ainda é aplicado $h$ vezes, mas como os elementos possuem dimensões menores, cada uma é $\frac{1}{h}$ vezes mais rápida, o que torna a escalabilidade do SDPA e MHA igual.

De forma prática, aplicar MHA ainda será mais lento que aplicar SDPA devido à projeção linear $W^O$. Porém, o tempo adicionado não aumenta em relação a nenhuma variável.

n_heads = 8
split_embed_dim = embed_dim // n_heads
outputs_projection = nn.Linear(embed_dim, embed_dim, bias=False)


def split_embeddings(
    embeddings: torch.Tensor,
    batch_size: int,
    n_tokens: int,
) -> torch.Tensor:
    splitted = embeddings.view(batch_size, n_tokens, n_heads, split_embed_dim)
    splitted = splitted.transpose(1, 2)

    return splitted


def join_embeddings(
    embeddings: torch.Tensor,
    batch_size: int,
    n_tokens: int,
) -> torch.Tensor:
    joined = embeddings.transpose(1, 2)
    joined = joined.contiguous()
    joined = joined.view(batch_size, n_tokens, embed_dim)
    return joined


def apply_multihead_attention(
    queries: torch.Tensor,
    keys: torch.Tensor,
    values: torch.Tensor,
) -> torch.Tensor:
    batch_size, n_tokens, _ = queries.size()

    queries = queries_projection(queries)
    keys = keys_projection(keys)
    values = values_projection(values)

    queries = split_embeddings(queries, batch_size, n_tokens)
    keys = split_embeddings(keys, batch_size, n_tokens)
    values = split_embeddings(values, batch_size, n_tokens)

    keys = keys.transpose(2, 3)
    scores = queries @ keys / (split_embed_dim**0.5)
    weights = F.softmax(scores, dim=3)

    outputs = weights @ values
    outputs = join_embeddings(outputs, batch_size, n_tokens)
    outputs = outputs_projection(outputs)

    return outputs

class MultiheadAttention(nn.Module):
    def __init__(self: Self, embed_dim: int, n_heads: int) -> None:
        super().__init__()

        self.embed_dim = embed_dim
        self.n_heads = n_heads
        self.split_embed_dim = self.embed_dim // self.n_heads

        self.queries_projection = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.keys_projection = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.values_projection = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
        self.outputs_projection = nn.Linear(self.embed_dim, self.embed_dim, bias=False)

    def split_embeddings(
        self: Self,
        embeddings: torch.Tensor,
        batch_size: int,
        n_tokens: int,
    ) -> torch.Tensor:
        splitted = embeddings.view(
            batch_size,
            n_tokens,
            self.n_heads,
            self.split_embed_dim,
        )

        splitted = splitted.transpose(1, 2)

        return splitted

    def join_embeddings(
        self: Self,
        embeddings: torch.Tensor,
        batch_size: int,
        n_tokens: int,
    ) -> torch.Tensor:
        joined = embeddings.transpose(1, 2)
        joined = joined.contiguous()
        joined = joined.view(batch_size, n_tokens, self.embed_dim)

        return joined

    def forward(
        self: Self,
        queries: torch.Tensor,
        keys: torch.Tensor,
        values: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, n_tokens, _ = queries.size()

        queries = self.queries_projection(queries)
        keys = self.keys_projection(keys)
        values = self.values_projection(values)

        queries = self.split_embeddings(queries, batch_size, n_tokens)
        keys = self.split_embeddings(keys, batch_size, n_tokens)
        values = self.split_embeddings(values, batch_size, n_tokens)

        keys = keys.transpose(2, 3)

        scores = queries @ keys / (self.split_embed_dim**0.5)
        weights = F.softmax(scores, dim=3)

        outputs = weights @ values
        outputs = self.join_embeddings(outputs, batch_size, n_tokens)
        outputs = self.outputs_projection(outputs)

        return outputs

Positional Encoding (PE)

Nenhum dos mecanismos de atenção explicados considera a posição dos elementos na sequência recebida na sequência gerada. Isso significa que alterar a ordem dos elementos não afeta o resultado, um efeito indesejado e que pode enviesar modelos durante o treinamento.

Antes da aplicação de qualquer mecanismo de atenção, as posições dos elementos originais são representadas usando uma técnica de PE, que gera uma sequência de embeddings onde cada elemento representa uma posição na sequência original de alguma forma.

Esses embeddings serão combinados com a sequência recebida para alterar a sequência recebida de forma que cada elemento tenha sua posição codificada individualmente.

A técnica de PE introduzida na arquitetura dos Transformers é baseada na função a seguir:

Onde:

• $p$ é a posição de um elemento na sequência recebida.
• $j$ é a posição de um item em um elemento da sequência recebida.
• $0 \leq i < \frac{d}{2}$, e $i$ é incrementado a cada dois itens consecutivos em um elemento da sequência recebida.
• $\theta$ é um hiperparâmetro.