scaled_dot_product_attention: atenția din Transformer în PyTorch: sintaxă Python, tensori și exemple practice

F.scaled_dot_product_attention calculează atenția din Transformer: scorurile q @ k.transpose(-2, -1) se înmulțesc cu factorul de scalare (implicit 1/sqrt(E), unde E este ultima dimensiune a lui query), se adaugă masca, se aplică softmax și dropout, iar rezultatul se înmulțește cu value. Documentația dă codul Python echivalent. Formele sunt (N, ..., Hq, L, E) pentru query și (N, ..., H, S, E) pentru key; rezultatul are forma (N, ..., Hq, L, Ev). Funcția este marcată ca beta.

Există trei implementări: FlashAttention-2, Memory-Efficient Attention și o implementare C++ care urmează formula. Pe backend-ul CUDA funcția poate alege kernelurile optimizate, după tipul și forma intrărilor; pe celelalte backend-uri se folosește implementarea PyTorch. Context managerul torch.nn.attention.sdpa_kernel() restrânge implementările permise; dacă cea cerută nu poate rula, apare un avertisment cu motivul.

Masca booleană are sensul invers față de key_padding_mask din nn.MultiheadAttention: aici True înseamnă că elementul participă la atenție. O mască de tip float se adună la scoruri. is_causal=True aplică masca triunghiulară inferioară; dacă dai și attn_mask, primești eroare. Dropout-ul se aplică mereu după dropout_p, deci la evaluare trebuie să dai 0.0.

Concepte cheie și caracteristici

Concept / Parametru	Descriere practică
`query, key, value`	Tensori `(N, ..., Hq, L, E)`, `(N, ..., H, S, E)`, `(N, ..., H, S, Ev)`.
`attn_mask`	Mască booleană (`True` = participă) sau float (adunată la scoruri), broadcastabilă la `(N, ..., Hq, L, S)`. Implicit `None`.
`dropout_p`	Probabilitatea de dropout pe ponderile de atenție. Implicit `0.0`; se aplică și la evaluare dacă e mai mare ca 0.
`is_causal`	Mască cauzală (triunghiulară inferioară). Nu se combină cu `attn_mask`. Implicit `False`.
`scale`	Factorul de scalare, doar ca argument cu nume. Implicit `1/sqrt(E)`.
`enable_gqa`	Grouped Query Attention (experimental): număr de capete query multiplu de numărul de capete key/value. Implicit `False`.
`sdpa_kernel(SDPBackend...)`	Context manager din `torch.nn.attention` care alege implementările permise (de exemplu `SDPBackend.FLASH_ATTENTION`, `SDPBackend.MATH`).

Exemple practice și scenarii reale

Exemplul 1: Self-attention cauzal cu mai multe capete

Dropout-ul se dezactivează în afara antrenării, cum recomandă documentația.

import torch
import torch.nn as nn
import torch.nn.functional as F

class AtentieCauzala(nn.Module):
    def __init__(self, dim=256, capete=8, p=0.1):
        super().__init__()
        self.capete, self.p = capete, p
        self.qkv = nn.Linear(dim, 3 * dim)
        self.proiectie = nn.Linear(dim, dim)

    def forward(self, x):                      # x: [N, L, dim]
        N, L, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=-1)
        q, k, v = (t.view(N, L, self.capete, -1).transpose(1, 2) for t in (q, k, v))
        y = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.p if self.training else 0.0,
        )
        y = y.transpose(1, 2).reshape(N, L, D)
        return self.proiectie(y)

print(AtentieCauzala()(torch.randn(2, 16, 256)).shape)   # torch.Size([2, 16, 256])

Exemplul 2: Mască de padding (True = token real)

q = k = v = torch.randn(2, 4, 6, 32)                  # [N, H, L, E]
lungimi = torch.tensor([6, 3])                        # al doilea exemplu are 3 tokeni reali
real = torch.arange(6)[None, :] < lungimi[:, None]    # [2, 6]
masca = real[:, None, None, :]                        # [2, 1, 1, 6] -> broadcast la [2, 4, 6, 6]
y = F.scaled_dot_product_attention(q, k, v, attn_mask=masca)

Exemplul 3: Forțarea unei implementări

Exemplul din documentație cere CUDA și float16; dacă FlashAttention nu acceptă intrările, apare un avertisment cu motivul.

from torch.nn.attention import sdpa_kernel, SDPBackend

q = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
k = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
v = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
    y = F.scaled_dot_product_attention(q, k, v)

Greșeli frecvente și capcane

Să lași dropout_p diferit de 0 la evaluare. Funcția nu știe dacă modulul e în train() sau eval().
Să muți o mască de la nn.MultiheadAttention fără s-o inversezi. Acolo True înseamnă „ignoră”, aici „participă” (~masca sau masca.logical_not()).
Să dai și is_causal=True, și attn_mask. Funcția aruncă eroare; pune masca cauzală direct în attn_mask dacă ai nevoie de amândouă.

`scaled_dot_product_attention: atenția din Transformer` în PyTorch & Deep Learning Practic