Aprenda Como Implementar Atenção Multi-Cabeça | Construindo Componentes de Transformers

Deslize para mostrar o menu


              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)

A atenção multi-cabeça funciona aplicando o mesmo mecanismo de self-attention de forma independente em várias cabeças em paralelo. Cada cabeça utiliza seu próprio conjunto de matrizes de pesos para queries, keys e values (W_q, W_k e W_v), projetando a entrada em subespaços de menor dimensão. A entrada x neste exemplo contém três tokens, cada um com tamanho de embedding 6, e é dividida em duas cabeças.

Para cada cabeça, o código repete as etapas padrão da self-attention: calcula as matrizes de query (Q), key (K) e value (V), calcula os scores de atenção realizando o produto escalar entre Q e K.T (com escala apropriada) e utiliza uma função softmax para obter os pesos de atenção. Esses pesos são então usados para combinar os vetores de value, produzindo a saída de cada cabeça.

Não há um novo mecanismo aqui, a atenção multi-cabeça é simplesmente o processo de self-attention realizado em paralelo para cada cabeça. Após o processamento de todas as cabeças, suas saídas são concatenadas ao longo da última dimensão. Essa abordagem permite que o modelo capture diferentes relações nos dados, já que cada cabeça pode focar em aspectos distintos da entrada. Os resultados impressos ilustram como executar o mesmo mecanismo de self-attention múltiplas vezes e combinar os resultados fornece uma representação mais rica para cada token.

Tudo estava claro?

Obrigado pelo seu feedback!

Seção 2. Capítulo 2

Pergunte à IA

Pergunte o que quiser ou experimente uma das perguntas sugeridas para iniciar nosso bate-papo

Seção 2. Capítulo 2