Introdução

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 gerou diversas variações que alcançaram o estado da arte em várias 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, são variações dos Transformers (como os Generative Pretrained Transformers ou 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.

Quando eu estava aprendendo como esses modelos funcionavam, encontrei diversos bons recursos explicando a arquitetura. No entanto, senti falta de uma explicação conceitual e prática de ponta a ponta, que permitisse compreender o funcionamento desses modelos a partir de uma única fonte. Quando precisei aprender sobre o tema, tive que recorrer a várias referências simultaneamente para obter uma visão abrangente, o que é especialmente desafiador para quem busca um aprendizado mais prático e deseja criar uma versão funcional do modelo. Além disso, embora existam excelentes materiais em inglês e outros idiomas, praticamente não há recursos em português que abordem Transformers no nível de profundidade que considero ideal. Acredito que isso seja um fator limitante para explicar o assunto no Brasil e em outros países lusófonos.

Por isso, meu objetivo é apresentar uma explicação de ponta a ponta sobre o funcionamento dos Transformers, abordando o contexto em que surgiram, o funcionamento conceitual de seus principais componentes e a arquitetura desses modelos. Vou ilustrar todos esses conceitos com exemplos práticos e implementações funcionais em PyTorch, de forma que você possa criar um modelo funcional usando apenas a teoria e o código apresentados aqui.

Pré-requisitos

Infelizmente, se eu não assumir absolutamente nenhum pré-requisito, o conteúdo ficará extenso demais para ser feito de uma vez. Pretendo escrever outros guias explicando esses pré-requisitos no futuro, e vou listar alguns recursos que você pode usar para estudar esses assuntos antes. Porém, para entender tudo que vou explicar aqui, é importante que você compreenda esses tópicos dos conteúdos a seguir:

Álgebra linear

Vamos trabalhar com conceitos como vetores, matrizes, tensores, produto escalar, produto vetorial e transposição de matrizes. Eu aprendi álgebra linear pelo livro "Álgebra Linear" de J.L. Boldrini na faculdade, mas hoje recomendo buscar conteúdos em português, como as aulas da UNIVESP ou da Khan Academy, que são ótimas para revisar ou aprender do zero.

Python e PyTorch

Você vai precisar conhecer o básico de Python: laços, estruturas de dados, condicionais, funções e um pouco de orientação a objetos (como criar classes, objetos e usar módulos). Todo o código do guia será feito em PyTorch, então é importante entender os conceitos principais do framework, como manipulação de tensores, uso de dispositivos, autograd e o submódulo torch.nn.

Se você nunca usou PyTorch antes, recomendo começar pelo tutorial oficial, que explica bem os conceitos básicos do framework. Se já tem alguma experiência e quer revisar, o curso aberto de Deep Learning da University of Amsterdam (UvA) tem um tutorial que resume de forma prática como programar redes neurais em PyTorch.

Deep Learning

Como transformers são redes neurais, vamos usar todos os conceitos envolvendo redes feedforward ou perceptrons multicamada (MLPs), como o uso do gradiente descendente, funções de ativação e backpropagation.

O curso Neural Networks and Deep Learning na Coursera é completo e muito acessível. Em português, além do livro "Dive into Deep Learning", sugiro os capítulos 7 e 8 do livro Aprendizado de Máquina - Uma Abordagem Estatística, que cobrem redes neurais.

Conclusão

Como o guia traz bastante código, você pode acessar tudo de forma separada na biblioteca language-models. Todo o código das funções e classes apresentadas aqui está organizado lá. O projeto é open source e pode ser instalado como uma biblioteca, embora ainda não esteja disponível no PyPI.

E, com isso, concluímos a apresentação. No próximo post, vamos iniciar o guia trazendo uma introdução sobre os Transformers e falando do seu contexto histórico.

Introdução

Nesse post, será feita uma introdução aos Transformers, explicando os problemas que estamos tentando resolver com essa arquitetura, introduzindo o contexto histórico que levou à sua criação e os desafios presentes na criação desse tipo de modelo.

Objetivo

De forma geral, Transformers são utilizados para resolver problemas de transdução de sequências, ou seja, tarefas em que é necessário gerar uma nova sequência de elementos a partir de uma sequência recebida. Essa abordagem permite que diversas aplicações sejam modeladas como transdução de sequências, incluindo tradução automática, geração de texto, sumarização e até mesmo síntese de moléculas.

Formalmente, um modelo de transdução de sequências estabelece uma relação entre uma sequência ordenada s=⟨s₁,s₂,⋯,sₙ⟩, composta por elementos de um conjunto enumerável S, e outra sequência t=⟨t₁,t₂,⋯,tₘ⟩, composta por elementos de um conjunto enumerável T.

Para facilitar a compreensão do funcionamento dos Transformers, muitos exemplos ao longo deste guia serão apresentados no contexto de uma aplicação específica. Considerando a relevância dos Transformers para o avanço de modelos de linguagem, como os Large Language Models, os exemplos deste guia focarão na construção de um modelo de linguagem.

Um pouco de história

Durante muitos anos, criar modelos eficientes para transdução de sequências foi um desafio em aberto para a comunidade científica. O principal obstáculo era superar limitações presentes nas arquiteturas anteriores aos Transformers.

Antes dos Transformers, as técnicas baseadas em redes neurais recorrentes (RNNs) apresentavam o melhor desempenho em tarefas como tradução automática, sendo a base da arquitetura utilizada pelo Google Tradutor em 2014. O processo de treinamento das RNNs enfrentava alguns problemas:

1. A natureza recursiva das RNNs pode causar o problema de gradient vanishing, em que os gradientes calculados durante o backpropagation se tornam tão pequenos que o modelo não consegue aprender de forma eficiente.
2. Como algumas operações precisam ser realizadas de forma sequencial, a inferência nesses modelos pode ser muito lenta, tornando o treinamento e o uso prático inviáveis em muitos casos.
3. Sem o uso de atenção, as RNNs têm dificuldade para capturar o significado de um trecho considerando o contexto completo da frase, o que pode prejudicar a qualidade do texto gerado.

O grande diferencial dos Transformers é serem baseados exclusivamente em redes neurais feedforward e mecanismos de atenção. Isso permite resolver os problemas mencionados e alcançar desempenho superior em tarefas de transdução de sequências.

Essa abordagem justifica o nome do artigo: do ponto de vista arquitetural, não são necessárias redes neurais recorrentes para criar modelos eficientes. Basta o uso de mecanismos de atenção, ou seja, "Attention is All you Need".

Atenção

Neste post, vou explicar o que é atenção e como funcionam os mecanismos apresentados no artigo "Attention is All You Need". No caminho, vamos ver como funcionam os mecanismos Scaled Dot-Product Attention e Multihead Attention, e ao final, também vamos implementar esses mecanismos do zero em PyTorch com foco em eficiência computacional.

O que é atenção?

No contexto de redes neurais para transdução de sequências, atenção refere-se à capacidade do modelo de considerar o contexto de cada elemento da sequência ao gerar um novo valor para cada elemento de entrada. Ou seja, o modelo pode "prestar atenção" em diferentes partes da sequência para produzir saídas mais precisas e contextualizadas.

Nesse processo, os pesos de atenção formam uma sequência intermediária que indica o quanto cada elemento da entrada deve ser considerado na geração de cada novo elemento. Ao aplicar esses pesos sobre os elementos originais, é gerada uma nova sequência de elementos para representar o conteúdo da sequência recebida, dessa vez representando melhor o contexto de cada elemento.

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

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 "trem".

Self-Attention

Todos os mecanismos utilizados na arquitetura dos Transformers funcionam combinando os próprios elementos da sequência recebida entre si para calcular a atenção de cada elemento. Ou seja, a atenção atribuída a cada elemento da sequência é baseada apenas nos elementos da sequência recebida, sem depender de informações externas.

Queries, Keys e Values

De forma simplificada, um mecanismo de atenção pode ser comparado a um dicionário em Python: chaves (keys ou K) são associadas a valores (values ou V), e é possível recuperar um valor a partir de uma consulta (query ou Q). No contexto do mecanismo de atenção, os papéis de query, key e value são desempenhados pelos próprios elementos da sequência recebida:

• Query: O elemento da sequência para o qual queremos representar de outra forma.
• Keys: Todos os elementos da sequência recebida, que funcionam como possíveis referências para determinar o contexto da query.
• Value: Um novo elemento que representa o significado da query, calculado a partir do elemento original e das keys.

Mecanismos de atenção

Dot-Product Attention (DPA)

Nesse mecanismo, os pesos de atenção de cada elemento são calculados usando o produto interno (dot product) entre as queries e as keys correspondentes.

Scaled Dot-Product Attention (SDPA)

Essa variação do DPA inclui um termo de normalização nos pesos para estabilizar os gradientes durante o backpropagation. A normalização pelo fator √d tem base empírica.

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 maneiras para melhorar a performance dos Transformers é aplicar projeções lineares separadas às queries, keys e values antes do seu uso nos mecanismos de atenção, multiplicando cada uma por uma matriz de parâmetros treinável específica (Wq, WK e WV).

Multihead Attention (MHA)

Essa variação do SDPA aplica o mecanismo de atenção h vezes em paralelo, cada uma com diferentes projeções lineares dos elementos da sequência. O número de cabeças h é um hiperparâmetro do modelo.

Como os pesos no SDPA são normalizados, cada cabeça de atenção não pode "prestar atenção" igualmente em todos os elementos ao mesmo tempo. O objetivo do uso de MHA é permitir que cada cabeça foque em diferentes padrões ou relações no contexto, enriquecendo a representação aprendida pelo modelo durante o treinamento.

class MultiheadAttention(nn.Module):
    def __init__(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)

Positional Encoding

No post anterior, vimos como funcionam mecanismos de atenção e concluímos que existem algumas limitações no seu uso direto. Nesse post, explicarei como funciona o Positional Encoding, uma técnica complementar à atenção para representar a posição dos elementos na sequência recebida.

Positional Encoding (PE)

Os mecanismos de atenção apresentados até aqui não levam em conta a posição dos elementos na sequência recebida ao gerar a sequência de saída. Isso significa que, se a ordem dos elementos for alterada, o resultado permanecerá o mesmo. Esse efeito é indesejado e pode introduzir vieses no modelo durante o treinamento.

Antes de aplicar qualquer mecanismo de atenção, é necessário representar na sequência recebida a posição de cada elemento de alguma forma. Na arquitetura, essa representação é feita usando Positional Encoding (PE), técnica que gera uma sequência de embeddings especiais: cada um representa a posição de um elemento na sequência, permitindo que o modelo entenda a ordem dos elementos a partir apenas dos elementos da sequência recebida.

Onde:

• p é a posição de um elemento na sequência recebida.
• j é a posição de um item escalar em um elemento da sequência recebida.
• i é um índice auxiliar tal que 0 ≤ i < d/2, e i é incrementado a cada dois itens consecutivos de j.
• θ é um hiperparâmetro que representa a escala das posições.

class PositionalEncoder(nn.Module):
    def __init__(self, config: PositionalEncoderConfig) -> None:
        super().__init__()
        self.config = config

    @torch.no_grad()
    def forward(self, embeddings: torch.Tensor) -> torch.Tensor:
        batch_size, n_tokens, _ = embeddings.size()
        
        indexes = torch.arange(self.config.embed_dim, dtype=torch.float)
        positions = torch.arange(n_tokens, dtype=torch.float)
        positions = positions.view(n_tokens, 1)

        i = torch.arange(self.config.embed_dim // 2)
        i = i.float()
        i = i.repeat_interleave(2)

        encodings = positions / (self.config.theta ** (2 * i / self.config.embed_dim))
        
        return embeddings + encodings

Modelos Autoregressivos

Nesse post, antes de entendermos a estrutura de um Transformer, vamos entender primeiro como funcionam os modelos autoregressivos. Eles são uma classe de modelos que recebem sequências de dados e geram novas sequências, um elemento por vez.

Modelos autoregressivos

Transformers são modelos que realizam transdução de sequências utilizando um processo autoregressivo. Essa característica define como os diferentes componentes do modelo são organizados e combinados ao longo da arquitetura.

Quando um modelo de transdução de sequências é autoregressivo, ele é treinado para modificar a sequência recebida por meio de um processo chamado shift, no qual:

• O primeiro elemento é removido.
• Todos os elementos são deslocados uma posição para trás.
• A última posição é preenchida com um novo elemento gerado pelo modelo.

O modelo gera um novo elemento, que é adicionado ao início da sequência gerada. Em seguida, a sequência pós-shift se torna a nova sequência recebida, e o processo se repete: a cada iteração, um novo elemento é gerado e acrescentado à sequência, até que algum critério de parada seja atingido.

Os critérios de parada podem ser:

1. Definir um número máximo de iterações.
2. Encerrar o algoritmo quando o último elemento gerado for igual a um valor especial, como o token <eos>.

Essa limitação é feita nos Transformers com o uso de uma attention mask, que zera parte dos pesos para garantir que um novo elemento não seja baseado nos valores dos elementos seguintes.

def attn_mask_like(size: tuple[int]) -> torch.Tensor:
    mask = torch.ones(size)
    mask = mask.triu(diagonal=1)
    mask = mask.bool()
    return mask

Conclusão

Usando um processo autogressivo, é possível criar modelos capazes de gerar sequências inteiras de elementos de forma iterativa. Esse é o segredo para criar inteligências artificiais capazes de traduzir textos, responder perguntas, e outras aplicações.

Agora temos todas as peças para montar nosso próprio modelo desse tipo. No próximo post, vamos ver como foi feita a arquitetura do primeiro Transformer, como descrito no artigo original.