ROFORMER: FÖRBÄTTRAD TRANSFORMATOR MED ROTERANDE POSITIONSBÄDDNING

Transformatorbaserade modeller är kända för sin förmåga att tolka och tolka komplex text. De förlitar sig på att förstå ordens ordning och sammanhang - uppgifter där traditionella positionskodningsmetoder har visat sina gränser. ROFORMER-modellen, som drivs av Rotary Position Embedding (RoPE), omdefinierar vårt tillvägagångssätt för positionell kodning för att lösa detta gap.

Traditionell positionskodning

Transformatorer behandlar text som en serie tokens och tillåter parallell bearbetning av sekvenser för större effektivitet. Men denna styrka medförde sin utmaning: modellens inneboende agnosticism till symbolisk ordning. Positionell kodning var svaret, tillhandahöll varje token en unik signatur som anger dess sekvensposition.

Absolute Position Inbäddningar

Till en början använde modeller som BERT absoluta positionsinbäddningar, som tilldelade en fast vektor till varje position i en sekvens. Även om den här metoden är enkel, saknar den i sig förmågan att anpassa sig till variationer i sekvenslängd eller att betona de relativa avstånden mellan tokens**, vilket är avgörande för att förstå många språkliga konstruktioner.

Relativ position Inbäddningar

För att fånga språkets dynamiska natur introducerades relativa positionsinbäddningar, med fokus på avståndet mellan tokens snarare än deras absoluta positioner. Trots sina konceptuella fördelar introducerade dessa inbäddningar beräkningskomplexitet, och misslyckades med att sömlöst integreras i Transformerss självuppmärksamhetsmekanism, vilket begränsade deras effektivitet.

ROFORMER och Rotary Position Inbäddning

Genom att inse begränsningarna hos befintliga positionskodningsstrategier introducerar ROFORMER Rotary Position Embedding (RoPE), ett tillvägagångssätt som kombinerar fördelarna med absolut och relativ positionsinformation utan deras respektive nackdelar.

Rotary Position Inbäddning

RoPE kodar positionsinformation med hjälp av rotationsmatriser, vilket gör att modellen inte bara förstår var en token är, utan hur den relaterar till alla andra token i en sekvens.

Credit: Original Paper

Den fungerar genom en geometrisk lins och behandlar tokenpositioner som punkter i ett flerdimensionellt utrymme som roteras för att markera deras sekventiella relationer. Denna rotation tillåter modellen att bevara och utnyttja både absoluta och relativa positionssignaler inom sin självuppmärksamhetsmekanism.

Implementering av RoPE

Implementering av RoPE innebär att koda varje tokens position till en rotationsmatris och applicera denna matris inom transformatorns självuppmärksamhetsmekanism. Denna process möjliggör en flexibel, dynamisk tolkning av positionsinformation, tar emot varierande sekvenslängder och fångar essensen av token-interrelationer utan betydande beräkningskostnader.

Först behöver du en funktion för att generera de roterande inbäddningarna, och sedan ska du integrera dessa inbäddningar i din modell. Exemplet nedan förutsätter att du är bekant med att skapa anpassade lager i Keras.

Steg 1: Definiera Rotary Embedding-funktionen

Denna funktion genererar de roterande inbäddningarna givet den maximala sekvenslängden och dimensionaliteten hos inbäddningarna.

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))

Denna linje beräknar inversen av exponentiellt skalade frekvenser baserat på positionsindex. Dessa frekvenser används för att generera sinusformade mönster för roterande inbäddningar, vilket hjälper till att koda den relativa positionsinformationen i sekvenser. Denna mekanism är särskilt användbar i uppgifter där förståelse av ordningen och den relativa placeringen av element är avgörande, till exempel vid naturlig språkbehandling eller tidsserieanalys.

I detalj:

tf.range(0, dim, 2, dtype=tf.float32) skapar ett värdeintervall från 0 upp till dim (exklusivt), steg med 2. Argumentet dtype=tf.float32 anger att elementen i denna tensor är 32-bitars flyttal. Om "dim" är 8, till exempel, skulle detta producera "[0, 2, 4, 6]".
Tensorn som produceras av tf.range divideras sedan med dimensionaliteten (dim) för inbäddningarna. Denna operation skalar ner dessa index till ett intervall mellan 0 och 1 (exklusivt om 'dim' är jämnt, något inklusive om 'dim' är udda, eftersom intervallsteget hoppar över vartannat värde). Om du fortsätter med exemplet med "dim" = 8, dividerat med 8 ger "[0,0, 0,25, 0,5, 0,75]".
Operationen 10000 ** (...) höjer 10 000 till styrkan för varje element i den tidigare skalade tensorn. Basen på 10 000 är något godtycklig, men är vald för att säkerställa att frekvenserna varierar över ett brett område, vilket hjälper modellen att mer effektivt skilja mellan olika positioner. För [0.0, 0.25, 0.5, 0.75] skulle den applicera effektdriften på var och en, vilket resulterar i värden mycket större för högre element.
Slutligen erhålls den inversa frekvensen genom att ta den reciproka (1/x) av värdena från föregående steg. De inversa frekvenserna är mindre för högre index, vilket betyder att element längre fram i sekvensen kommer att ha mindre frekvenser, vilket påverkar hur deras positioner kodas in i modellen. Detta gör att inbäddningarna kan skalas på ett sätt där relativa positioner kan härledas genom modellens uppmärksamhetsmekanismer.

Linjen:

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

använder TensorFlows tf.einsum-funktion, ett verktyg som möjliggör kortfattade och effektiva uttryck av tensoroperationer med Einsteins summationsnotation.

Denna operation beräknar effektivt den yttre produkten av vektorerna "t" och "inv_freq", vilket resulterar i en matris där varje element "(i, j)" är produkten av "i"-te elementet av "t" och "t" j-te elementet i inv_freq`. Denna matris ("freqs") representerar de frekvenser som används för att generera de sinusformade mönstren för de roterande inbäddningarna.

Steg 2: Anpassat Keras-lager för roterande inbäddningar

Låt oss nu skapa ett anpassat Keras-lager som applicerar roterande inbäddningar på ingångstensorn. Detta lager antar att ingångstensorn har formen "(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

Raden inbäddningar = self.rotary_embeddings[:seq_len] väljer lämplig delmängd av förberäknade roterande inbäddningar baserat på den aktuella inmatningssekvenslängden. Eftersom längden på sekvenser kan variera från en batch till en annan säkerställer denna skivningsoperation att endast de inbäddningar som motsvarar den faktiska sekvenslängden används.

Variabeln 'inbäddningar' har nu en tensor av formen '(seq_len, embedding_dim)', där 'seq_len' är längden på sekvenserna i den aktuella batchen och 'embedding_dim' är dimensionaliteten för inbäddningarna. Denna tensor innehåller roterande positionsinbäddningar för varje position i sekvensen upp till "seq_len".

emb = tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) kombinerar sinus- och cosinustransformationer av positionsfrekvenser till en enda tensor:

-tf.cos(freqs) och tf.sin(freqs) tillämpar cosinus- respektive sinustransformationerna på freqs-tensorn. "freqs"-tensorn innehåller frekvensvärden för varje position i inmatningssekvensen och varje dimension av inbäddningsutrymmet, beräknat baserat på sekvenspositionerna och de omvända frekvenserna för inbäddningsdimensionerna. Sinus- och cosinusfunktionerna appliceras elementvis, vilket resulterar i två tensorer med samma form som "freqs". Dessa transformationer hjälper till att koda positionen på ett sätt som fångar den cykliska karaktären hos positionsförhållanden, vilket underlättar modellens förmåga att förstå relativa positioner.

-tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) sammanfogar de cosinus- och sinustransformerade tensorerna längs den sista axeln (betecknad med axis=-1). Att sammanfoga dessa tensorer sida vid sida fördubblar effektivt dimensionaliteten för "freqs"-tensorn, där den första halvan representerar cosinustransformerade värden och den andra halvan representerar sinustransformerade värden för varje position. Sammankopplingen säkerställer att varje positionskodning innehåller både sinus- och cosinusinformation, vilket möjliggör bevarande av information om både amplituden och fasen för positionssignalerna.

Den sammanlänkade tensor "emb" håller nu de fullständiga roterande inbäddningarna för ingångspositionerna. Formen på "emb" kommer att vara densamma som "freqs" i sina två första dimensioner (motsvarande sekvenspositioner och inbäddningsdimensioner), men dess sista dimension kommer att vara dubbelt så stor, vilket tar hänsyn till både sinus- och cosinusvärden. Dessa inbäddningar används för att modulera inbäddningsinbäddningarna genom att lägga till positionsinformation på ett rotationsmässigt ekvivariant sätt.

-cos_emb = inbäddningar[:, Inga, :self.dim // 2]:

Det första kolonet : betyder "välj alla element i denna dimension", vilket i det här fallet hänvisar till alla positioner i sekvensen.
Ingen används för att lägga till en extra dimension, vilket gör tensorn tredimensionell. Detta görs ofta för att säkerställa kompatibilitet med vissa operationer som förväntar sig indata av ett specifikt antal dimensioner. Till exempel, när du utför elementvis multiplikation med en annan tensor som är 3-dimensionell, måste formerna anpassas enligt sändningsregler.
:self.dim // 2, väljer den första halvan av dimensionerna i den sista axeln. Eftersom "inbäddningsdimensionen" fördubblas för att inkludera både sinus- och cosinusvärden, väljs enbart cosinuskomponenterna för inbäddningarna genom att dividera med 2.

Steg 3: Integration med en Keras-modell

Efter att ha definierat RotaryEmbeddingLayer kan du integrera det i din Keras-modell. Detta lager bör appliceras på dina inbäddningar innan du matar in dem i uppmärksamhetslager eller några efterföljande modelllager.

Här är ett förenklat exempel på hur man integrerar de roterande inbäddningarna i en modell:

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()