RoFormer: Pinahusay na Transformer na May Rotary Position Embedding

Mga modelong nakabatay sa transformer ay sikat sa kanilang kakayahang mag-parse at mag-interpret ng kumplikadong teksto. Umaasa sila sa pag-unawa sa pagkakasunud-sunod at konteksto ng mga salita - mga gawain kung saan ipinakita ng tradisyonal na positional encoding na mga pamamaraan ang kanilang mga limitasyon. Sa pagtugon sa gap na ito, ang modelong ROFORMER, na pinapagana ng Rotary Position Embedding (RoPE), ay muling tumutukoy sa aming diskarte sa positional encoding.

Tradisyunal na Positional Encoding

Itinuturing ng mga transformer ang text bilang isang serye ng mga token, at pinapayagan ang parallel processing ng mga sequence para sa higit na kahusayan. Gayunpaman, dinala ng lakas na ito ang hamon nito: ang likas na agnostisismo ng modelo sa pagkakasunud-sunod ng token. Positional encoding ang sagot, nagbibigay sa bawat token ng natatanging pirma na nagsasaad ng sequence position nito.

Mga Pag-embed ng Ganap na Posisyon

Sa una, ang mga modelong tulad ng BERT ay gumamit ng mga pag-embed ng ganap na posisyon, na nagtatalaga ng isang nakapirming vector sa bawat posisyon sa isang pagkakasunud-sunod. Ang pamamaraang ito, bagama't diretso, ay likas na walang kakayahang umangkop sa mga pagkakaiba-iba ng haba ng pagkakasunud-sunod, o upang bigyang-diin ang mga kaugnay na distansya sa pagitan ng mga token, na kritikal para sa pag-unawa sa maraming linguistic na konstruksyon.

Mga Pag-embed ng Kamag-anak na Posisyon

Upang makuha ang dynamic na kalikasan ng wika, ipinakilala ang mga relatibong pag-embed ng posisyon, na tumutuon sa distansya sa pagitan ng mga token kaysa sa kanilang mga ganap na posisyon. Sa kabila ng kanilang konseptong kalamangan, ang mga pag-embed na ito ay nagpakilala ng computational complexity, at nabigong maayos na maisama sa mekanismo ng self-attention ng Transformers, na nililimitahan ang kanilang bisa.

ROFORMER at Rotary Position Embedding

Kinikilala ang mga limitasyon ng umiiral na mga diskarte sa positional encoding, ipinakilala ng ROFORMER ang Rotary Position Embedding (RoPE), isang diskarte na pinagsasama ang mga benepisyo ng absolute at relative position information nang walang kani-kanilang mga disbentaha.

Rotary Position Embedding

Ang RoPE ay nag-e-encode ng positional na impormasyon gamit ang rotation matrice, na nagbibigay-daan sa modelo na maunawaan hindi lamang kung nasaan ang isang token, ngunit kung paano ito nauugnay sa bawat iba pang token sa isang sequence.

Credit: ArXiv

Gumagana ito sa pamamagitan ng isang geometric na lens, tinatrato ang mga posisyon ng token bilang mga punto sa isang multi-dimensional na espasyo na iniikot upang markahan ang kanilang mga magkakasunod na relasyon. Ang pag-ikot na ito ay nagbibigay-daan sa modelo na mapanatili at mapagsamantalahan ang parehong absolute at relatibong positional na mga pahiwatig sa loob ng mekanismo ng sarili nitong atensyon.

Pagpapatupad ng RoPE

Ang pagpapatupad ng RoPE ay nagsasangkot ng pag-encode ng bawat posisyon ng token sa isang rotation matrix, at paglalapat ng matrix na ito sa loob ng self-attention mechanism ng Transformer. Ang prosesong ito ay nagbibigay-daan para sa isang flexible, dynamic na interpretasyon ng positional na impormasyon, pagtanggap ng iba't ibang haba ng sequence at pagkuha ng esensya ng mga token interrelationship na walang makabuluhang computational overhead.

Una, kakailanganin mo ng isang function upang makabuo ng mga rotary embeddings, at pagkatapos ay isasama mo ang mga embedding na ito sa iyong modelo. Ipinapalagay ng halimbawa sa ibaba na pamilyar ka sa paggawa ng mga custom na layer sa Keras.

Hakbang 1: Tukuyin ang Rotary Embedding Function

Binubuo ng function na ito ang mga rotary embeddings na binibigyan ng maximum na haba ng pagkakasunud-sunod, at ang dimensionality ng mga 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))

Kinakalkula ng linyang ito ang kabaligtaran ng exponentially scaled frequency batay sa mga indeks ng posisyon. Ang mga frequency na ito ay ginagamit sa pagbuo ng mga sinusoidal pattern para sa mga rotary embeddings, na tumutulong sa pag-encode ng relatibong positional na impormasyon sa mga sequence. Ang mekanismong ito ay partikular na kapaki-pakinabang sa mga gawain kung saan ang pag-unawa sa pagkakasunud-sunod at relatibong pagpoposisyon ng mga elemento ay mahalaga, tulad ng sa natural na pagpoproseso ng wika o pagsusuri ng time series.

Sa mga detalye:

Ang tf.range(0, dim, 2, dtype=tf.float32) ay lumilikha ng hanay ng mga value na nagsisimula mula 0 hanggang dim (eksklusibo), humahakbang ng 2. Tinutukoy ng argumento ng dtype=tf.float32 na ang mga elemento ng tensor na ito ay 32-bit na mga floating-point na numero. Kung ang dim ay 8, halimbawa, ito ay magbubunga ng [0, 2, 4, 6].
Ang tensor na ginawa ng tf.range ay hinati sa dimensionality (dim) ng mga embed. Pinababa ng operasyong ito ang mga indeks na ito sa isang hanay sa pagitan ng 0 at 1 (eksklusibo kung ang dim ay pantay, bahagyang kasama kung ang dilim ay kakaiba, dahil nilalaktawan ng hakbang ng hanay ang bawat iba pang value). Ang pagpapatuloy ng halimbawa na may dilam = 8, hinahati sa 8 ay [0.0, 0.25, 0.5, 0.75].
Ang 10000 ** (...) na operasyon ay nagtataas ng 10,000 sa kapangyarihan ng bawat elemento sa dating naka-scale na tensor. Ang base ng 10,000 ay medyo arbitrary, ngunit pinili upang matiyak na ang mga frequency ay nag-iiba-iba sa isang malawak na hanay, na tumutulong sa modelo na makilala ang iba't ibang mga posisyon nang mas epektibo. Para sa [0.0, 0.25, 0.5, 0.75], ilalapat nito ang power operation sa bawat isa, na magreresulta sa mga value na mas malaki para sa mas matataas na elemento.
Panghuli, ang inverse frequency ay nakuha sa pamamagitan ng pagkuha ng reciprocal (1/x) ng mga value mula sa nakaraang hakbang. Ang mga inverse frequency ay mas maliit para sa mas matataas na mga indeks, ibig sabihin, ang mga elemento pa sa sequence ay magkakaroon ng mas maliliit na frequency, na makakaapekto sa kung paano naka-encode ang kanilang mga posisyon sa modelo. Nagbibigay-daan ito sa mga pag-embed na mag-scale sa paraang kung saan maaaring mahinuha ang mga kamag-anak na posisyon sa pamamagitan ng mga mekanismo ng atensyon ng modelo.

Ang linya:

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

gumagamit ng tf.einsum function ng TensorFlow, isang tool na nagbibigay-daan para sa maikli at mahusay na pagpapahayag ng mga operasyon ng tensor gamit ang Einstein summation notation.

Ang operasyong ito ay epektibong kinakalkula ang panlabas na produkto ng t at inv_freq na mga vector, na nagreresulta sa isang matrix kung saan ang bawat elemento (i, j) ay produkto ng i-th elemento ng t at ang j-ika elemento ng inv_freq. Ang matrix na ito (freqs) ay kumakatawan sa mga frequency na ginagamit upang bumuo ng mga sinusoidal na pattern para sa mga rotary embeddings.

Hakbang 2: Custom Keras Layer para sa Rotary Embeddings

Ngayon, gumawa tayo ng custom na Keras layer na naglalapat ng mga rotary embeddings sa input tensor. Ipinapalagay ng layer na ito na ang input tensor ay may hugis (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

Pinipili ng linyang embeddings = self.rotary_embeddings[:seq_len] ang naaangkop na subset ng paunang na-compute na rotary embeddings batay sa kasalukuyang haba ng sequence ng input. Dahil ang haba ng mga pagkakasunud-sunod ay maaaring mag-iba mula sa isang batch patungo sa isa pa, tinitiyak ng slicing operation na ito na ang mga embeddings na tumutugma sa aktwal na haba ng sequence ang ginagamit.

Ang variable na embeddings ay mayroon na ngayong tensor ng hugis (seq_len, embedding_dim), kung saan ang seq_len ay ang haba ng mga sequence sa kasalukuyang batch, at ang embedding_dim ay ang dimensionality ng mga embedding. Ang tensor na ito ay naglalaman ng mga rotary positional embeddings para sa bawat posisyon sa sequence hanggang sa seq_len.

Pinagsasama ng emb = tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) ang mga pagbabago sa sine at cosine ng mga positional na frequency sa isang tensor:

-tf.cos(freqs) at tf.sin(freqs) ilapat ang cosine at sine transformations, ayon sa pagkakabanggit, sa freqs tensor. Ang freqs tensor ay naglalaman ng mga frequency value para sa bawat posisyon sa input sequence at bawat dimensyon ng embedding space, na kinakalkula batay sa mga sequence positions at ang inverse frequency ng mga dimensyon ng pag-embed. Ang mga function ng sine at cosine ay inilalapat sa element-wise, na nagreresulta sa dalawang tensor ng parehong hugis bilang freqs. Nakakatulong ang mga pagbabagong ito sa pag-encode ng posisyon sa paraang kumukuha ng paikot na katangian ng mga positional na relasyon, na nagpapadali sa kakayahan ng modelo na maunawaan ang mga kaugnay na posisyon.

-tf.concat((tf.cos(freqs), tf.sin(freqs)), axis=-1) pinagsasama-sama ang cosine at sine transformed tensors kasama ang huling axis (na tinutukoy ng axis=-1). Ang pagsasama-sama ng mga tensor na ito nang magkatabi ay epektibong nagdodoble sa dimensionality ng freqs tensor, kung saan ang unang kalahati ay kumakatawan sa mga value na na-transform ng cosine at ang pangalawang kalahati ay kumakatawan sa mga halaga ng sine-transformed para sa bawat posisyon. Tinitiyak ng concatenation na ang bawat positional encoding ay naglalaman ng parehong sine at cosine na impormasyon, na nagbibigay-daan sa pagpapanatili ng impormasyon tungkol sa parehong amplitude at phase ng positional signal.

Ang pinagsama-samang tensor na emb ay hawak na ngayon ang kumpletong rotary embeddings para sa mga posisyon ng input. Ang hugis ng emb ay magiging kapareho ng freqs sa unang dalawang dimensyon nito (naaayon sa mga posisyon ng pagkakasunud-sunod at mga dimensyon ng pag-embed), ngunit ang huling dimensyon nito ay magiging dalawang beses na mas malaki, na isinasaalang-alang ang parehong mga halaga ng sine at cosine. Ang mga pag-embed na ito ay ginagamit upang i-modulate ang mga input embeddings sa pamamagitan ng pagdaragdag ng positional na impormasyon sa isang rotationally equivariant na paraan.

-cos_emb = mga pag-embed[:, Wala, :self.dim // 2]:

Ang unang colon : ay nangangahulugang "piliin ang lahat ng elemento sa dimensyong ito," na, sa kasong ito, ay tumutukoy sa lahat ng mga posisyon sa pagkakasunud-sunod.
Ginagamit ang Wala upang magdagdag ng karagdagang dimensyon, na ginagawang 3-dimensional ang tensor. Madalas itong ginagawa upang matiyak ang pagiging tugma sa ilang partikular na operasyon na umaasa sa mga input ng isang partikular na bilang ng mga dimensyon. Halimbawa, kapag nagsasagawa ng element-wise multiplication sa isa pang tensor na 3-dimensional, ang mga hugis ay dapat na nakahanay ayon sa mga panuntunan sa pagsasahimpapawid.
:self.dim // 2, pinipili ang unang kalahati ng mga dimensyon sa huling axis. Dahil ang embedding_dimension ay dinoble upang isama ang parehong mga halaga ng sine at cosine, ang paghahati sa 2 ay epektibong pinipili lamang ang mga bahagi ng cosine ng mga pag-embed.

Hakbang 3: Pagsasama sa isang Keras Model

Pagkatapos tukuyin ang RotaryEmbeddingLayer, maaari mo itong isama sa iyong modelo ng Keras. Dapat ilapat ang layer na ito sa iyong mga embed bago ilagay ang mga ito sa mga layer ng atensyon o anumang kasunod na mga layer ng modelo.

Narito ang isang pinasimpleng halimbawa kung paano isama ang mga rotary embeddings sa isang 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()