RoFormer: Transformador mellorado con incorporación de posición rotativa

Modelos baseados en transformadores son famosos pola súa capacidade para analizar e interpretar textos complexos. Confían en comprender a orde e o contexto das palabras, tarefas nas que os métodos tradicionais de codificación posicional mostraron os seus límites. Para abordar esta lagoa, o modelo ROFORMER, impulsado pola Incrustación de posición rotativa (RoPE), redefine o noso enfoque da codificación posicional.

Codificación posicional tradicional

Os transformadores tratan o texto como unha serie de tokens e permiten o procesamento paralelo de secuencias para unha maior eficiencia. Non obstante, esta fortaleza trouxo o seu desafío: o agnosticismo inherente do modelo á orde de símbolos. Codificación de posición foi a resposta, proporcionando a cada token unha sinatura única que indica a súa posición de secuencia.

Incrustacións de posición absoluta

Inicialmente, modelos como BERT utilizaban insercións de posición absoluta, asignando un vector fixo a cada posición nunha secuencia. Este método, aínda que sinxelo, inherentemente carece da capacidade de adaptarse ás variacións de lonxitude das secuencias ou de enfatizar as distancias relativas entre as fichas, fundamental para comprender moitos construtos lingüísticos.

Incrustacións de posición relativa

Para capturar a natureza dinámica da linguaxe, introducíronse incrustacións de posicións relativas, centrándose na distancia entre as fichas e non nas súas posicións absolutas. A pesar da súa vantaxe conceptual, estas incorporacións introduciron complexidade computacional e non se integraron perfectamente no mecanismo de autoatención de Transformers, limitando a súa eficacia.

ROFORMER e incrustación de posicións rotativas

Recoñecendo as limitacións das estratexias de codificación posicional existentes, ROFORMER presenta Rotary Position Embedding (RoPE), un enfoque que combina os beneficios da información de posición absoluta e relativa sen os seus respectivos inconvenientes.

Incorporación de posicións rotativas

RoPE codifica a información posicional mediante matriz de rotación, o que permite que o modelo comprenda non só onde está un token, senón como se relaciona con todos os outros tokens nunha secuencia.

Credit: ArXiv

Opera a través dunha lente xeométrica, tratando as posicións dos tokens como puntos nun espazo multidimensional que se xiran para marcar as súas relacións secuenciais. Esta rotación permite que o modelo preserve e explote indicios posicionais tanto absolutos como relativos dentro do seu mecanismo de autoatención.

Implementación de RoPE

Implementar RoPE implica codificar a posición de cada ficha nunha matriz de rotación e aplicar esta matriz dentro do mecanismo de autoatención do Transformador. Este proceso permite unha interpretación flexible e dinámica da información posicional, acomodando lonxitudes de secuencia variables e capturando a esencia das interrelacións de tokens sen sobrecarga computacional significativa.

Primeiro, necesitarás unha función para xerar as incorporacións rotativas e, a continuación, integrarás estas incorporacións no teu modelo. O seguinte exemplo asume que estás familiarizado coa creación de capas personalizadas en Keras.

Paso 1: Defina a función de incrustación rotativa

Esta función xera as incrustacións rotativas dada a lonxitude máxima da secuencia e a dimensionalidade das incrustacións.

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 liña calcula a inversa das frecuencias escaladas exponencialmente en función dos índices de posición. Estas frecuencias úsanse para xerar patróns sinusoidais para incrustacións rotativas, o que axuda a codificar a información de posición relativa en secuencias. Este mecanismo é particularmente útil en tarefas nas que a comprensión da orde e o posicionamento relativo dos elementos é crucial, como no procesamento da linguaxe natural ou a análise de series temporais.

En detalles:

tf.range(0, dim, 2, dtype=tf.float32) crea un rango de valores a partir de 0 ata dim (exclusivo), pasando por 2. O argumento dtype=tf.float32 especifica que os elementos deste tensor son números de coma flotante de 32 bits. Se dim é 8, por exemplo, isto produciría [0, 2, 4, 6].
O tensor producido por tf.range divídese entón pola dimensionalidade (dim) das incorporacións. Esta operación reduce estes índices a un rango entre 0 e 1 (exclusivo se dim é par, lixeiramente inclusivo se dim é impar, porque o paso de rango omite todos os demais valores). Continuando o exemplo con dim = 8, dividindo por 8 obtense [0,0, 0,25, 0,5, 0,75].
A operación 10.000 ** (...) eleva 10.000 á potencia de cada elemento no tensor escalado previamente. A base de 10.000 é algo arbitraria, pero elíxese para garantir que as frecuencias varían nun amplo rango, o que axuda ao modelo a diferenciar as diferentes posicións de forma máis eficaz. Para [0.0, 0.25, 0.5, 0.75], aplicaría a operación de potencia a cada un, resultando en valores moito maiores para elementos superiores.
Finalmente, a frecuencia inversa obtense tomando o recíproco (1/x) dos valores do paso anterior. As frecuencias inversas son máis pequenas para índices máis altos, o que significa que os elementos máis adiante na secuencia terán frecuencias máis pequenas, o que afectará a forma en que se codifican as súas posicións no modelo. Isto permite que as incrustacións escalan de forma que se poidan inferir posicións relativas a través dos mecanismos de atención do modelo.

A liña:

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

usa a función tf.einsum de TensorFlow, unha ferramenta que permite a expresión concisa e eficiente de operacións de tensor usando a notación de suma de Einstein.

Esta operación calcula eficazmente o produto exterior dos vectores t e inv_freq, resultando nunha 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 frecuencias que se usan para xerar os patróns sinusoidais para as incrustacións rotativas.

Paso 2: Capa Keras personalizada para incorporacións rotativas

Agora, imos crear unha capa Keras personalizada que aplique incorporacións rotativas ao tensor de entrada. Esta capa asume que o tensor de entrada ten 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 liña embeddings = self.rotary_embeddings[:seq_len] selecciona o subconxunto axeitado de incrustacións rotativas precalculadas en función da lonxitude da secuencia de entrada actual. Dado que a lonxitude das secuencias pode variar dun lote a outro, esta operación de corte garante que só se utilicen as incorporacións correspondentes á lonxitude real da secuencia.

A variable embeddings agora contén un tensor de forma (seq_len, embedding_dim), onde seq_len é a lonxitude das secuencias do lote actual e embedding_dim é a dimensionalidade das incorporacións. Este tensor contén as incrustacións posicionais rotativas para cada posición na secuencia ata seq_len.

emb = tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) combina transformacións de seno e coseno de frecuencias posicionais nun único tensor:

-tf.cos(freqs) e tf.sin(freqs) aplican as transformacións coseno e seno, respectivamente, ao tensor freqs. O tensor freqs contén valores de frecuencia para cada posición na secuencia de entrada e cada dimensión do espazo de incrustación, calculados en función das posicións da secuencia e das frecuencias inversas das dimensións de incrustación. As funcións seno e coseno aplícanse por elementos, dando lugar a dous tensores da mesma forma que freqs. Estas transformacións axudan a codificar a posición dun xeito que capte a natureza cíclica das relacións posicionais, facilitando a capacidade do modelo para comprender as posicións relativas.

-tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) concatena os tensores transformados coseno e seno ao longo do último eixe (indicado por axis=-1). A concatenación destes tensores lado a lado duplica efectivamente a dimensionalidade do tensor "freqs", coa primeira metade representando valores transformados en coseno e a segunda metade representando valores transformados en seno para cada posición. A concatenación garante que cada codificación posicional conteña información de seno e coseno, o que permite conservar a información tanto da amplitude como da fase dos sinais posicionais.

O tensor concatenado emb agora contén as incrustacións rotativas completas para as posicións de entrada. A forma de emb será a mesma que freqs nas súas dúas primeiras dimensións (correspondentes ás posicións da secuencia e ás dimensións de incrustación), pero a súa última dimensión será o dobre, contando os valores de seno e coseno. Estas incorporacións utilízanse para modular as incorporacións de entrada engadindo información posicional de forma rotacionalmente equivalente.

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

Os primeiros dous puntos : significan "seleccionar todos os elementos desta dimensión", que, neste caso, refírese a todas as posicións da secuencia.
None úsase para engadir unha dimensión adicional, facendo que o tensor sexa tridimensional. Isto adoita facerse para garantir a compatibilidade con determinadas operacións que esperan entradas dun número específico de dimensións. Por exemplo, cando se realiza a multiplicación por elementos con outro tensor tridimensional, as formas deben aliñarse segundo as regras de transmisión.
:self.dim // 2, selecciona a primeira metade das dimensións no último eixe. Dado que a embedding_dimension duplícase para incluír valores de seno e coseno, dividindo por 2 só se seleccionan os compoñentes coseno das incorporacións.

Paso 3: Integración cun modelo Keras

Despois de definir o RotaryEmbeddingLayer, podes integralo no teu modelo Keras. Esta capa debe aplicarse ás súas incrustacións antes de alimentalas en capas de atención ou en capas de modelos posteriores.

Aquí tes un exemplo simplificado de como integrar as incrustacións rotativas nun 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()