Lära Hur man implementerar multi-head attention

Svep för att visa menyn


              12345678910111213141516171819202122232425262728293031323334353637383940
            
import numpy as np

# Toy input: 3 tokens, embedding size 6
x = np.array([
    [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
    [0.6, 0.5, 0.4, 0.3, 0.2, 0.1],
    [0.2, 0.1, 0.4, 0.3, 0.6, 0.5],
])

# Multi-head attention parameters
num_heads = 2
embed_size = x.shape[1]
head_dim = embed_size // num_heads

# Random weights for queries, keys, and values for each head
np.random.seed(42)
W_q = np.random.randn(num_heads, embed_size, head_dim)
W_k = np.random.randn(num_heads, embed_size, head_dim)
W_v = np.random.randn(num_heads, embed_size, head_dim)

attention_outputs = []

for head in range(num_heads):
    # Linear projections
    Q = x @ W_q[head]
    K = x @ W_k[head]
    V = x @ W_v[head]
    
    # Scaled dot-product attention
    scores = Q @ K.T / np.sqrt(head_dim)
    attention_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    head_output = attention_weights @ V
    attention_outputs.append(head_output)

# Concatenate outputs from all heads
multi_head_output = np.concatenate(attention_outputs, axis=-1)

print("Output from head 0:\n", attention_outputs[0])
print("Output from head 1:\n", attention_outputs[1])
print("Concatenated multi-head output:\n", multi_head_output)

Multi-head attention fungerar genom att samma self-attention-mekanism tillämpas oberoende över flera huvuden parallellt. Varje huvud använder sina egna viktmatriser för queries, keys och values (W_q, W_k och W_v), och projicerar indata till lägre-dimensionella delrymden. Indatan x i detta exempel innehåller tre token, var och en med en inbäddningsstorlek på 6, och delas upp i två huvuden.

För varje huvud upprepas de vanliga self-attention-stegen: query (Q), key (K) och value (V) matriser beräknas, attention-poäng räknas ut genom att ta skalärprodukten av Q och K.T (med lämplig skalning), och en softmax-funktion används för att erhålla attention-vikter. Dessa vikter används sedan för att kombinera value-vektorerna, vilket ger varje huvuds utdata.

Det finns ingen ny mekanism här, multi-head attention är helt enkelt self-attention-processen utförd parallellt för varje huvud. När alla huvuden har bearbetats, sammanfogas deras utdata längs den sista dimensionen. Detta tillvägagångssätt gör det möjligt för modellen att fånga olika relationer i datan, eftersom varje huvud kan fokusera på olika aspekter av indatan. De utskrivna resultaten illustrerar hur samma self-attention-mekanism, körd flera gånger och kombinerad, ger en rikare representation för varje token.

Var allt tydligt?

Tack för dina kommentarer!

Avsnitt 2. Kapitel 2

Fråga AI

Fråga vad du vill eller prova någon av de föreslagna frågorna för att starta vårt samtal

Avsnitt 2. Kapitel 2