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-sekvenca 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 obraća pažnju na sve druge tokene, 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

Samo pomnoži svaku dimenzionalnu vrednost upita sa relevantnom vrednošću svakog tokena i saberi rezultate. Dobijaš 1 vrednost po paru tokena.

Za svaku reč u rečenici, izračunaj 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 gubi se u matematičkim terminima, cilj ove funkcije je jednostavan, normalizuj sve težine tako da ukupno sumiraju 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 sabiraju na 1.

Izračunavanje eksponencijala:

Izračunavanje sume:

Izračunavanje težina pažnje:

Korak 3: Izračunavanje vektora konteksta

Samo uzmi svaku težinu pažnje i pomnoži je sa dimenzijama relevantnog tokena, a zatim saberi sve dimenzije da dobiješ 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"**:

Sabiranje 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čunaj rezultate pažnje: Koristi skalarni proizvod između umetanja ciljne reči i umetanja svih reči u sekvenci.

2. Normalizuj rezultate da bi dobio težine pažnje: Primeni softmax funkciju na rezultate pažnje da dobiješ težine koje se sabiraju na 1.

3. Izračunaj vektor konteksta: Pomnoži umetanje svake reči sa njenom težinom pažnje i saberi 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čunaj upite, ključeve 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)

Inicijalizuj 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: Izračunavanje skalirane dot-produkta pažnje

Izračunavanje ocena pažnje

Slično prethodnom primeru, ali ovoga puta, umesto korišćenja vrednosti dimenzija tokena, koristimo ključnu matricu tokena (već izračunatu korišćenjem dimenzija):. Dakle, za svaku upit qi​ i ključ kj​:

Skaliranje ocena

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

Ocena se deli kvadratnim korenom dimenzija jer dot proizvodi mogu postati veoma veliki i ovo pomaže u regulisanju.

Primena Softmax-a za dobijanje težina pažnje: Kao u inicijalnom primeru, normalizujte sve vrednosti tako da se saberu na 1.

Step 3: Izračunavanje kontekstualnih vektora

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

Primer koda

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 nasumičnim vrednostima, koristi nn.Linear da označi sve težine kao parametre 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 sumiraju 1. Ova maska postavlja rezultate pažnje budućih tokena na negativnu beskonačnost, osiguravajući da nakon softmax-a, njihova težina pažnje bude 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

1. 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 učenje, možemo primeniti dropout na težine pažnje nakon softmax operacije. Dropout nasumično 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 u skladu sa brojem 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 kao prednost 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