4. Attention Mechanisms

Mehanizmi pažnje i samopažnja u neuronskim mrežama

Mehanizmi pažnje omogućavaju neuronskim mrežama da fokusiraju na specifične delove ulaza prilikom generisanja svakog dela izlaza. Dodeljuju različite težine različitim ulazima, pomažući modelu da odluči koji su ulazi najrelevantniji za zadatak. Ovo je ključno u zadacima poput mašinskog prevođenja, gde je razumevanje konteksta cele rečenice neophodno za tačan prevod.

Cilj ove četvrte faze je vrlo jednostavan: Primeni neke mehanizme pažnje. Ovo će biti mnogo ponovljenih slojeva koji će uhvatiti odnos reči u rečniku sa njenim susedima u trenutnoj rečenici koja se koristi za obuku LLM. Za ovo se koristi mnogo slojeva, tako da će mnogo parametara za obuku hvatati ove informacije.

Razumevanje mehanizama pažnje

U tradicionalnim modelima sekvenca-za-sekvencu koji se koriste za prevođenje jezika, model kodira ulaznu sekvencu u vektorski kontekst fiksne veličine. Međutim, ovaj pristup se suočava sa problemima sa dugim rečenicama jer fiksni vektorski kontekst možda neće uhvatiti sve potrebne informacije. Mehanizmi pažnje rešavaju ovo ograničenje omogućavajući modelu da razmatra sve ulazne tokene prilikom generisanja svakog izlaznog tokena.

Primer: Mašinsko prevođenje

Razmotrite prevođenje nemačke rečenice "Kannst du mir helfen diesen Satz zu übersetzen" na engleski. Prevod reč po reč ne bi proizveo gramatički ispravnu englesku rečenicu zbog razlika u gramatičkim strukturama između jezika. Mehanizam pažnje omogućava modelu da se fokusira na relevantne delove ulazne rečenice prilikom generisanja svake reči izlazne rečenice, što dovodi do tačnijeg i koherentnijeg prevoda.

Uvod u samopažnju

Samopažnja, ili intra-pažnja, je mehanizam gde se pažnja primenjuje unutar jedne sekvence kako bi se izračunala reprezentacija te sekvence. Omogućava svakom tokenu u sekvenci da se obraća svim drugim tokenima, pomažući modelu da uhvati zavisnosti između tokena bez obzira na njihovu udaljenost u sekvenci.

Ključni koncepti

• Tokeni: Pojedinačni elementi ulazne sekvence (npr. reči u rečenici).

• Umetanja: Vektorske reprezentacije tokena, koje hvataju semantičke informacije.

• Težine pažnje: Vrednosti koje određuju važnost svakog tokena u odnosu na druge.

Izračunavanje težina pažnje: Primer korak po korak

Razmotrimo rečenicu "Hello shiny sun!" i predstavimo svaku reč sa 3-dimenzionalnim umetanjima:

• Hello: [0.34, 0.22, 0.54]

• shiny: [0.53, 0.34, 0.98]

• sun: [0.29, 0.54, 0.93]

Naš cilj je da izračunamo vektor konteksta za reč "shiny" koristeći samopažnju.

Korak 1: Izračunavanje rezultata pažnje

Jednostavno pomnožite svaku dimenzionalnu vrednost upita sa relevantnom vrednošću svakog tokena i saberite rezultate. Dobijate 1 vrednost po paru tokena.

Za svaku reč u rečenici, izračunajte rezultat pažnje u odnosu na "shiny" izračunavanjem skalarne produkcije njihovih umetanja.

Rezultat pažnje između "Hello" i "shiny"

Rezultat pažnje između "shiny" i "shiny"

Rezultat pažnje između "sun" i "shiny"

Korak 2: Normalizacija rezultata pažnje da bi se dobile težine pažnje

Ne gubite se u matematičkim terminima, cilj ove funkcije je jednostavan, normalizujte sve težine tako da se ukupno saberu na 1.

Pored toga, softmax funkcija se koristi jer naglašava razlike zbog eksponencijalnog dela, olakšavajući prepoznavanje korisnih vrednosti.

Primeni softmax funkciju na rezultate pažnje da bi ih pretvorio u težine pažnje koje se saberu na 1.

Izračunavanje eksponencijala:

Izračunavanje sume:

Izračunavanje težina pažnje:

Korak 3: Izračunavanje vektora konteksta

Jednostavno uzmite svaku težinu pažnje i pomnožite je sa dimenzijama relevantnog tokena, a zatim saberite sve dimenzije da dobijete samo 1 vektor (vektor konteksta)

Vektor konteksta se izračunava kao ponderisana suma umetanja svih reči, koristeći težine pažnje.

Izračunavanje svake komponente:

• Ponderisano umetanje "Hello":

* **Ponderisano umetanje "shiny"**:

* **Ponderisano umetanje "sun"**:

Saberanje ponderisanih umetanja:

vektor konteksta=[0.0779+0.2156+0.1057, 0.0504+0.1382+0.1972, 0.1237+0.3983+0.3390]=[0.3992,0.3858,0.8610]

Ovaj vektor konteksta predstavlja obogaćeno umetanje za reč "shiny," uključujući informacije iz svih reči u rečenici.

Sažetak procesa

1. Izračunajte rezultate pažnje: Koristite skalarni proizvod između umetanja ciljne reči i umetanja svih reči u sekvenci.

2. Normalizujte rezultate da biste dobili težine pažnje: Primeni softmax funkciju na rezultate pažnje da biste dobili težine koje se saberu na 1.

3. Izračunajte vektor konteksta: Pomnožite umetanje svake reči sa njenom težinom pažnje i saberite rezultate.

Samopažnja sa težinama koje se mogu obučavati

U praksi, mehanizmi samopažnje koriste težine koje se mogu obučavati da nauče najbolje reprezentacije za upite, ključeve i vrednosti. Ovo uključuje uvođenje tri matrice težina:

Upit je podatak koji se koristi kao i ranije, dok su matrice ključeva i vrednosti samo nasumične matrice koje se mogu obučavati.

Korak 1: Izračunavanje upita, ključeva i vrednosti

Svaki token će imati svoju matricu upita, ključeva i vrednosti množenjem svojih dimenzionalnih vrednosti sa definisanim matricama:

Ove matrice transformišu originalna umetanja u novi prostor pogodan za izračunavanje pažnje.

Primer

Pretpostavljajući:

• Ulazna dimenzija din=3 (veličina umetanja)

• Izlazna dimenzija dout=2 (željena dimenzija za upite, ključeve i vrednosti)

Inicijalizujte matrice težina:

import torch.nn as nn

d_in = 3
d_out = 2

W_query = nn.Parameter(torch.rand(d_in, d_out))
W_key = nn.Parameter(torch.rand(d_in, d_out))
W_value = nn.Parameter(torch.rand(d_in, d_out))

Izračunajte upite, ključeve i vrednosti:

queries = torch.matmul(inputs, W_query)
keys = torch.matmul(inputs, W_key)
values = torch.matmul(inputs, W_value)

Step 2: Compute Scaled Dot-Product Attention

Compute Attention Scores

Slično prethodnom primeru, ali ovoga puta, umesto da koristimo vrednosti dimenzija tokena, koristimo ključnu matricu tokena (već izračunatu koristeći dimenzije):. Dakle, za svaki upit qi​ i ključ kj​:

Scale the Scores

Da bismo sprečili da dot proizvodi postanu preveliki, skaliramo ih kvadratnim korenom dimenzije ključa dk​:

Rezultat se deli kvadratnim korenom dimenzija jer dot proizvodi mogu postati veoma veliki i ovo pomaže da se regulišu.

Apply Softmax to Obtain Attention Weights: Kao u inicijalnom primeru, normalizujte sve vrednosti tako da se saberu na 1.

Step 3: Compute Context Vectors

Kao u inicijalnom primeru, samo saberite sve matrice vrednosti množeći svaku sa njenom težinom pažnje:

Code Example

Uzimajući primer sa https://github.com/rasbt/LLMs-from-scratch/blob/main/ch03/01_main-chapter-code/ch03.ipynb možete proveriti ovu klasu koja implementira funkcionalnost samopaznje o kojoj smo govorili:

import torch

inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your     (x^1)
[0.55, 0.87, 0.66], # journey  (x^2)
[0.57, 0.85, 0.64], # starts   (x^3)
[0.22, 0.58, 0.33], # with     (x^4)
[0.77, 0.25, 0.10], # one      (x^5)
[0.05, 0.80, 0.55]] # step     (x^6)
)

import torch.nn as nn
class SelfAttention_v2(nn.Module):

def __init__(self, d_in, d_out, qkv_bias=False):
super().__init__()
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

def forward(self, x):
keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)

attn_scores = queries @ keys.T
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

context_vec = attn_weights @ values
return context_vec

d_in=3
d_out=2
torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

Napomena da se umesto inicijalizacije matrica sa slučajnim vrednostima, koristi nn.Linear da se označe sve težine kao parametri za obuku.

Uzročna Pažnja: Sakrivanje Budućih Reči

Za LLM-ove želimo da model uzima u obzir samo tokene koji se pojavljuju pre trenutne pozicije kako bi predvideo sledeći token. Uzročna pažnja, takođe poznata kao maskirana pažnja, to postiže modifikovanjem mehanizma pažnje kako bi se sprečio pristup budućim tokenima.

Primena Maski za Uzročnu Pažnju

Da bismo implementirali uzročnu pažnju, primenjujemo masku na rezultate pažnje pre softmax operacije tako da preostali rezultati i dalje zbrajaju 1. Ova maska postavlja rezultate pažnje budućih tokena na negativnu beskonačnost, osiguravajući da nakon softmax-a, njihove težine pažnje budu nula.

Koraci

1. Izračunaj Rezultate Pažnje: Isto kao i pre.

2. Primeni Masku: Koristi gornju trougaastu matricu ispunjenu negativnom beskonačnošću iznad dijagonale.

mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) * float('-inf')
masked_scores = attention_scores + mask

3. Primeni Softmax: Izračunaj težine pažnje koristeći maskirane rezultate.

attention_weights = torch.softmax(masked_scores, dim=-1)

Maskiranje Dodatnih Težina Pažnje sa Dropout-om

Da bismo sprečili prekomerno prilagođavanje, možemo primeniti dropout na težine pažnje nakon softmax operacije. Dropout slučajno postavlja neke od težina pažnje na nulu tokom obuke.

dropout = nn.Dropout(p=0.5)
attention_weights = dropout(attention_weights)

Redovni dropout je oko 10-20%.

Code Example

Code example from https://github.com/rasbt/LLMs-from-scratch/blob/main/ch03/01_main-chapter-code/ch03.ipynb:

import torch
import torch.nn as nn

inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your     (x^1)
[0.55, 0.87, 0.66], # journey  (x^2)
[0.57, 0.85, 0.64], # starts   (x^3)
[0.22, 0.58, 0.33], # with     (x^4)
[0.77, 0.25, 0.10], # one      (x^5)
[0.05, 0.80, 0.55]] # step     (x^6)
)

batch = torch.stack((inputs, inputs), dim=0)
print(batch.shape)

class CausalAttention(nn.Module):

def __init__(self, d_in, d_out, context_length,
dropout, qkv_bias=False):
super().__init__()
self.d_out = d_out
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New

def forward(self, x):
b, num_tokens, d_in = x.shape
# b is the num of batches
# num_tokens is the number of tokens per batch
# d_in is the dimensions er token

keys = self.W_key(x) # This generates the keys of the tokens
queries = self.W_query(x)
values = self.W_value(x)

attn_scores = queries @ keys.transpose(1, 2) # Moves the third dimension to the second one and the second one to the third one to be able to multiply
attn_scores.masked_fill_(  # New, _ ops are in-place
self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_size
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1
)
attn_weights = self.dropout(attn_weights)

context_vec = attn_weights @ values
return context_vec

torch.manual_seed(123)

context_length = batch.shape[1]
d_in = 3
d_out = 2
ca = CausalAttention(d_in, d_out, context_length, 0.0)

context_vecs = ca(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

Proširenje jedne glave pažnje na više glava pažnje

Višeglava pažnja u praktičnom smislu se sastoji od izvršavanja više instanci funkcije samopaznje, svaka sa svojim težinama, tako da se izračunavaju različiti konačni vektori.

Primer koda

Moguće je ponovo koristiti prethodni kod i samo dodati omotač koji ga pokreće nekoliko puta, ali ovo je optimizovana verzija sa https://github.com/rasbt/LLMs-from-scratch/blob/main/ch03/01_main-chapter-code/ch03.ipynb koja obrađuje sve glave u isto vreme (smanjujući broj skupih for petlji). Kao što možete videti u kodu, dimenzije svake oznake su podeljene u različite dimenzije prema broju glava. Na ovaj način, ako oznaka ima 8 dimenzija i želimo da koristimo 3 glave, dimenzije će biti podeljene u 2 niza od 4 dimenzije, a svaka glava će koristiti jedan od njih:

class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert (d_out % num_heads == 0), \
"d_out must be divisible by num_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer(
"mask",
torch.triu(torch.ones(context_length, context_length),
diagonal=1)
)

def forward(self, x):
b, num_tokens, d_in = x.shape
# b is the num of batches
# num_tokens is the number of tokens per batch
# d_in is the dimensions er token

keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection

return context_vec

torch.manual_seed(123)

batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)

context_vecs = mha(batch)

print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

Za još jednu kompaktno i efikasnu implementaciju možete koristiti torch.nn.MultiheadAttention klasu u PyTorch-u.

Kratak odgovor ChatGPT-a o tome zašto je bolje podeliti dimenzije tokena među glavama umesto da svaka glava proverava sve dimenzije svih tokena:

Iako bi omogućavanje svakoj glavi da obrađuje sve dimenzije ugrađivanja moglo izgledati korisno jer bi svaka glava imala pristup punim informacijama, standardna praksa je da se podele dimenzije ugrađivanja među glavama. Ovaj pristup balansira računarsku efikasnost sa performansama modela i podstiče svaku glavu da uči raznolike reprezentacije. Stoga je deljenje dimenzija ugrađivanja generalno poželjnije od toga da svaka glava proverava sve dimenzije.

References

• https://www.manning.com/books/build-a-large-language-model-from-scratch