Attention dai principi fondamentali

I transformer sono più facili da debuggare quando le shape dei tensori sono noiose. Parti da una sequenza di stati nascosti:

$$X \in \mathbb{R}^{T \times d_{model}}$$

Proiettala in query, key e value:

$$Q = XW_Q,\quad K = XW_K,\quad V = XW_V$$

La matrice degli score

Gli score di attention sono similarità a coppie tra ogni query e ogni key:

$$S = \frac{QK^\top}{\sqrt{d_k}}$$

Se $Q \in \mathbb{R}^{T \times d_k}$ e $K \in \mathbb{R}^{T \times d_k}$, allora:

$$S \in \mathbb{R}^{T \times T}$$

Quella matrice quadrata è tutto il trucco. La riga $i$ dice da quali posizioni il token $i$ dovrebbe leggere.

Una versione PyTorch minimale

import torch

def attention(q, k, v, mask=None):
    scale = q.size(-1) ** -0.5
    scores = q @ k.transpose(-2, -1) * scale

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    weights = scores.softmax(dim=-1)
    return weights @ v

Il bug comune

Il modo più facile per rompere attention è normalizzare sull'asse sbagliato:

# Sbagliato per la attention normale.
weights = scores.softmax(dim=-2)

# Corretto: ogni riga di query diventa una distribuzione sulle key.
weights = scores.softmax(dim=-1)

L'output mantiene la larghezza dei value:

$$\operatorname{Attention}(Q,K,V) = \operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \in \mathbb{R}^{T \times d_v}$$