ROFORMER: TRANSFORMADOR MELHORADO COM INCORPORAÇÃO EM POSIÇÃO ROTATIVA

Modelos baseados em transformadores são famosos por sua capacidade de analisar e interpretar textos complexos. Eles baseiam-se na compreensão da ordem e do contexto das palavras – tarefas nas quais os métodos tradicionais de codificação posicional mostraram os seus limites. Abordando essa lacuna, o modelo ROFORMER, desenvolvido com Rotary Position Embedding (RoPE), redefine nossa abordagem à codificação posicional.

Codificação posicional tradicional

Os transformadores tratam o texto como uma série de tokens e permitem o processamento paralelo de sequências para maior eficiência. No entanto, esta força trouxe o seu desafio: o agnosticismo inerente ao modelo em relação à ordem simbólica. Codificação posicional foi a resposta, fornecendo a cada token uma assinatura única denotando sua posição de sequência.

Incorporações de posição absoluta

Inicialmente, modelos como o BERT usavam incorporações de posição absoluta, atribuindo um vetor fixo a cada posição em uma sequência. Este método, embora direto, inerentemente não tem a capacidade de se adaptar às variações de comprimento de sequência ou de enfatizar as distâncias relativas entre os tokens, essenciais para a compreensão de muitas construções linguísticas.

Incorporações de posição relativa

Para capturar a natureza dinâmica da linguagem, foram introduzidas incorporações de posição relativa, concentrando-se na distância entre os tokens em vez de em suas posições absolutas. Apesar de sua vantagem conceitual, essas incorporações introduziram complexidade computacional e não conseguiram se integrar perfeitamente ao mecanismo de autoatenção dos Transformers, limitando sua eficácia.

ROFORMER e incorporação de posição rotativa

Reconhecendo as limitações das estratégias de codificação posicional existentes, a ROFORMER apresenta Rotary Position Embedding (RoPE), uma abordagem que combina os benefícios das informações de posição absoluta e relativa sem suas respectivas desvantagens.

Incorporação de posição rotativa

RoPE codifica informações posicionais usando matrizes de rotação, permitindo que o modelo entenda não apenas onde um token está, mas como ele se relaciona com todos os outros tokens em uma sequência.

Credit: Original Paper

Ele opera através de lentes geométricas, tratando as posições dos tokens como pontos em um espaço multidimensional que são girados para marcar suas relações sequenciais. Essa rotação permite que o modelo preserve e explore pistas posicionais absolutas e relativas dentro de seu mecanismo de autoatenção.

Implementando RoPE

A implementação do RoPE envolve codificar a posição de cada token em uma matriz de rotação e aplicar essa matriz dentro do mecanismo de autoatenção do Transformer. Este processo permite uma interpretação flexível e dinâmica de informações posicionais, acomodando comprimentos variados de sequência e capturando a essência dos inter-relacionamentos de tokens sem sobrecarga computacional significativa.

Primeiro, você precisará de uma função para gerar os embeddings rotativos e, em seguida, integrará esses embeddings ao seu modelo. O exemplo abaixo pressupõe que você esteja familiarizado com a criação de camadas personalizadas no Keras.

Etapa 1: Definir a função de incorporação rotativa

Esta função gera os embeddings rotativos dado o comprimento máximo da sequência e a dimensionalidade dos embeddings.

from tensorflow.keras.layers import Layer
import numpy as np

def get_rotary_embedding(dim, max_seq_len):
    inv_freq = 1.0 / (10000 ** (tf.range(0, dim, 2, dtype=tf.float32) / dim))
    t = tf.range(max_seq_len, dtype=tf.float32)
    freqs = tf.einsum('i,j->ij', t, inv_freq)
    emb = tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1)
    return emb

inv_freq = 1.0 / (10000 ** (tf.range(0, dim, 2, dtype=tf.float32) / dim))

Esta linha calcula o inverso das frequências em escala exponencial com base nos índices de posição. Essas frequências são utilizadas na geração de padrões senoidais para incorporações rotativas, o que auxilia na codificação das informações posicionais relativas em sequências. Este mecanismo é particularmente útil em tarefas onde a compreensão da ordem e do posicionamento relativo dos elementos é crucial, como no processamento de linguagem natural ou na análise de séries temporais.

Em detalhes:

tf.range(0, dim, 2, dtype=tf.float32) cria um intervalo de valores começando de 0 até dim (exclusivo), avançando em 2. O argumento dtype=tf.float32 especifica que os elementos deste tensor são números de ponto flutuante de 32 bits. Se dim for 8, por exemplo, isso produziria [0, 2, 4, 6].
O tensor produzido por tf.range é então dividido pela dimensionalidade (dim) dos embeddings. Esta operação reduz esses índices para um intervalo entre 0 e 1 (exclusivo se dim for par, ligeiramente inclusivo se dim for ímpar, porque a etapa do intervalo ignora todos os outros valores). Continuando o exemplo com dim = 8, dividir por 8 resulta em [0,0, 0,25, 0,5, 0,75].
A operação 10000 ** (...) eleva 10.000 à potência de cada elemento no tensor previamente escalado. A base de 10.000 é um tanto arbitrária, mas é escolhida para garantir que as frequências variem em uma ampla faixa, o que ajuda o modelo a diferenciar as diferentes posições de forma mais eficaz. Para [0,0, 0,25, 0,5, 0,75], aplicaria a operação de potência a cada um, resultando em valores muito maiores para elementos mais altos.
Por fim, a frequência inversa é obtida tomando o recíproco (1/x) dos valores do passo anterior. As frequências inversas são menores para índices mais altos, o que significa que os elementos mais adiante na sequência terão frequências menores, afetando a forma como suas posições são codificadas no modelo. Isso permite que os embeddings sejam dimensionados de forma que as posições relativas possam ser inferidas por meio dos mecanismos de atenção do modelo.

A linha:

freqs = tf.einsum('i,j->ij', t, inv_freq)

usa a função tf.einsum do TensorFlow, uma ferramenta que permite a expressão concisa e eficiente de operações tensorais usando a notação de soma de Einstein.

Esta operação calcula efetivamente o produto externo dos vetores t e inv_freq, resultando em uma matriz onde cada elemento (i, j) é o produto do i-ésimo elemento de t e o j-ésimo elemento de inv_freq. Esta matriz (freqs) representa as frequências que são usadas para gerar os padrões senoidais para os embeddings rotativos.

Etapa 2: Camada Keras personalizada para incorporações rotativas

Agora, vamos criar uma camada Keras personalizada que aplica incorporações rotativas ao tensor de entrada. Esta camada assume que o tensor de entrada tem a forma (batch_size, sequence_length, embedding_dim).

class RotaryEmbeddingLayer(Layer):
    def __init__(self, dim, max_seq_len, **kwargs):
        super().__init__(**kwargs)
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.rotary_embeddings = get_rotary_embedding(dim, max_seq_len)

    def call(self, inputs):
        seq_len = tf.shape(inputs)[1]
        embeddings = self.rotary_embeddings[:seq_len]
        
        cos_emb = embeddings[:, None, :self.dim // 2]
        sin_emb = embeddings[:, None, self.dim // 2:]
        
        # Decompose inputs into sine and cosine components
        inputs_cos = inputs[..., :self.dim // 2]
        inputs_sin = inputs[..., self.dim // 2:]
        
        # Apply rotary embeddings
        rotated_cos = inputs_cos * cos_emb - inputs_sin * sin_emb
        rotated_sin = inputs_sin * cos_emb + inputs_cos * sin_emb
        
        return tf.concat([rotated_cos, rotated_sin], axis=-1)

    def get_config(self):
        config = super().get_config()
        config.update({
            "dim": self.dim,
            "max_seq_len": self.max_seq_len
        })
        return config

A linha embeddings = self.rotary_embeddings[:seq_len] seleciona o subconjunto apropriado de embeddings rotativos pré-calculados com base no comprimento da sequência de entrada atual. Como o comprimento das sequências pode variar de um lote para outro, esta operação de fatiamento garante que apenas os embeddings correspondentes ao comprimento real da sequência sejam usados.

A variável embeddings agora contém um tensor de forma (seq_len, embedding_dim), onde seq_len é o comprimento das sequências no lote atual e embedding_dim é a dimensionalidade dos embeddings. Este tensor contém os embeddings posicionais rotativos para cada posição na sequência até seq_len.

emb = tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) combina transformações de seno e cosseno de frequências posicionais em um único tensor:

-tf.cos(freqs) e tf.sin(freqs) aplicam as transformações cosseno e seno, respectivamente, ao tensor freqs. O tensor freqs contém valores de frequência para cada posição na sequência de entrada e cada dimensão do espaço de incorporação, calculados com base nas posições da sequência e nas frequências inversas das dimensões de incorporação. As funções seno e cosseno são aplicadas elemento a elemento, resultando em dois tensores com a mesma forma de freqs. Estas transformações ajudam a codificar a posição de uma forma que capta a natureza cíclica das relações posicionais, facilitando a capacidade do modelo de compreender posições relativas.

-tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) concatena os tensores transformados cosseno e seno ao longo do último eixo (denotado por axis=-1). Concatenar esses tensores lado a lado efetivamente dobra a dimensionalidade do tensor freqs, com a primeira metade representando valores transformados em cosseno e a segunda metade representando valores transformados em seno para cada posição. A concatenação garante que cada codificação posicional contenha informações de seno e cosseno, o que permite a preservação de informações sobre a amplitude e a fase dos sinais posicionais.

O tensor concatenado emb agora contém os embeddings rotativos completos para as posições de entrada. A forma de emb será a mesma de freqs em suas duas primeiras dimensões (correspondendo às posições de sequência e dimensões de incorporação), mas sua última dimensão será duas vezes maior, contabilizando os valores de seno e cosseno. Esses embeddings são usados para modular os embeddings de entrada adicionando informações posicionais de maneira rotacionalmente equivariante.

-cos_emb = embeddings[:, Nenhum, :self.dim // 2]:

Os primeiros dois pontos : significa "selecionar todos os elementos nesta dimensão", que, neste caso, refere-se a todas as posições na sequência.
None é usado para adicionar uma dimensão adicional, tornando o tensor tridimensional. Isso geralmente é feito para garantir a compatibilidade com determinadas operações que esperam entradas de um número específico de dimensões. Por exemplo, ao realizar a multiplicação elemento a elemento com outro tensor tridimensional, as formas devem se alinhar de acordo com as regras de transmissão.
:self.dim // 2, seleciona a primeira metade das dimensões no último eixo. Como embedding_dimension é duplicado para incluir os valores de seno e cosseno, dividir por 2 seleciona efetivamente apenas os componentes de cosseno dos embeddings.

Etapa 3: Integração com um modelo Keras

Depois de definir RotaryEmbeddingLayer, você pode integrá-lo ao seu modelo Keras. Esta camada deve ser aplicada aos seus embeddings antes de alimentá-los nas camadas de atenção ou em qualquer camada de modelo subsequente.

Aqui está um exemplo simplificado de como integrar os embeddings rotativos em um modelo:

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Embedding, Dense

max_seq_len = 512
embedding_dim = 64

inp = Input(shape=(max_seq_len,))
x = Embedding(input_dim=10000, output_dim=embedding_dim)(inp)
x = RotaryEmbeddingLayer(dim=embedding_dim, max_seq_len=max_seq_len)(x)
# Add your model's layers here, e.g., Transformer blocks
x = Dense(1, activation='sigmoid')(x)

model = Model(inputs=inp, outputs=x)
model.summary()