Nyströmformer Approssimare l’auto-attenzione in tempo lineare e memoria mediante il metodo di Nyström

Nyströmformer approssima l'auto-attenzione in tempo lineare e memoria con il metodo di Nyström.

Introduzione

I trasformatori hanno mostrato prestazioni notevoli in vari compiti di elaborazione del linguaggio naturale e visione artificiale. Il loro successo può essere attribuito al meccanismo di auto-attenzione, che cattura le interazioni a coppie tra tutti i token in un input. Tuttavia, il meccanismo di auto-attenzione standard ha complessità di tempo e memoria di O ( n 2 ) O(n^2) O ( n 2 ) (dove n n n è la lunghezza della sequenza di input), rendendo costoso allenarsi su sequenze di input lunghe.

Il Nyströmformer è uno dei molti modelli efficienti dei trasformatori che approssima l’auto-attenzione standard con complessità O ( n ) O(n) O ( n ). Il Nyströmformer mostra prestazioni competitive in vari compiti di NLP e CV, migliorando l’efficienza dell’auto-attenzione standard. Lo scopo di questo post del blog è fornire ai lettori una panoramica del metodo Nyström e come può essere adattato per approssimare l’auto-attenzione.

Metodo Nyström per l’approssimazione di matrici

Alla base del Nyströmformer c’è il metodo Nyström per l’approssimazione di matrici. Ci consente di approssimare una matrice campionando alcune delle sue righe e colonne. Consideriamo una matrice P n × n P^{n \times n} P n × n , che è costosa da calcolare nella sua interezza. Quindi, invece, la approssimiamo usando il metodo Nyström. Iniziamo campionando m m m righe e colonne da P P P . Possiamo quindi disporre le righe campionate e le colonne come segue:

Rappresentazione di P come matrice a blocchi

• Presentazione del Private Hub Un nuovo modo di creare con l’apprendimento automatico
• Allenare e Ottimizzare i Modelli di Sentence Transformers
• Deployare 🤗 ViT su Kubernetes con TF Serving

Ora abbiamo quattro sottomatrici: A P , B P , F P , A_P, B_P, F_P, A P ​ , B P ​ , F P ​ , e C P C_P C P ​ , con dimensioni m × m , m × ( n − m ) , ( n − m ) × m m \times m, m \times (n – m), (n – m) \times m m × m , m × ( n − m ) , ( n − m ) × m e ( n − m ) × ( n − m ) (n – m) \times (n – m) ( n − m ) × ( n − m ) rispettivamente. Le m m m colonne campionate sono contenute in A P A_P A P ​ e F P F_P F P ​ , mentre le m m m righe campionate sono contenute in A P A_P A P ​ e B P B_P B P ​ . Quindi, le voci di A P , B P , A_P, B_P, A P ​ , B P ​ , e F P F_P F P ​ ci sono note, e stimiamo C P C_P C P ​ . Secondo il metodo Nyström, C P C_P C P ​ è dato da:

C P = F P A P + B P C_P = F_P A_P^+ B_P C P ​ = F P ​ A P + ​ B P ​

Qui, + + + indica la pseudoinversa di Moore-Penrose. Quindi, l’approssimazione di Nyström di P , P ^ P, \hat{P} P , P ^ può essere scritta come:

Approssimazione di Nyström di P

Come mostrato nella seconda riga, P ^ \hat{P} P ^ può essere espresso come prodotto di tre matrici. La ragione di ciò diventerà chiara in seguito.

Possiamo approssimare l’auto-attenzione con il metodo Nyström?

Il nostro obiettivo è alla fine approssimare la matrice softmax nell’auto-attenzione standard: S = softmax Q K T d \frac{QK^T}{\sqrt{d}} d ​ Q K T ​

Qui, Q Q Q e K K K indicano le query e le chiavi rispettivamente. Seguendo la procedura discussa in precedenza, campioneremmo m m m righe e colonne da S S S , formiamo quattro sottomatrici e otteniamo S ^ \hat{S} S ^ :

Approssimazione di Nyström di S

Ma, cosa significa campionare una colonna da S S S ? Significa selezionare un elemento da ciascuna riga. Ricordiamo come viene calcolato S: l’operazione finale è una softmax riga per riga. Per trovare una singola voce in una riga, dobbiamo accedere a tutte le altre voci (per il denominatore nella softmax). Quindi, campionare una colonna richiede di conoscere tutte le altre colonne nella matrice. Pertanto, non possiamo applicare direttamente il metodo Nyström per approssimare la matrice softmax.

Come possiamo adattare il metodo di Nyström per approssimare l’auto-attenzione?

Invece di campionare da S S S , gli autori propongono di campionare landmark (o punti di Nyström) da query e chiavi. Denotiamo i landmark query e i landmark chiave rispettivamente come Q ~ \tilde{Q} Q ~ ​ e K ~ \tilde{K} K ~ . Q ~ \tilde{Q} Q ~ ​ e K ~ \tilde{K} K ~ possono essere utilizzati per costruire tre matrici corrispondenti a quelle dell’approssimazione di Nyström di S S S . Definiamo le seguenti matrici:

F ~ = s o f t m a x ( Q K ~ T d ) A ~ = s o f t m a x ( Q ~ K ~ T d ) + B ~ = s o f t m a x ( Q ~ K T d ) \tilde{F} = softmax(\frac{Q\tilde{K}^T}{\sqrt{d}}) \hspace{40pt} \tilde{A} = softmax(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d}})^+ \hspace{40pt} \tilde{B} = softmax(\frac{\tilde{Q}K^T}{\sqrt{d}}) F ~ = s o f t m a x ( d ​ Q K ~ T ​ ) A ~ = s o f t m a x ( d ​ Q ~ ​ K ~ T ​ ) + B ~ = s o f t m a x ( d ​ Q ~ ​ K T ​ )

Le dimensioni di F ~ \tilde{F} F ~ , A ~ \tilde{A} A ~ , e B ~ ) a r e ( n × m , m × m , \tilde{B}) sono \\(n \times m, m \times m, B ~ ) a r e ( n × m , m × m , e m × n m \times n m × n rispettivamente. Sostituiamo le tre matrici nell’approssimazione di Nyström di S S S con le nuove matrici che abbiamo definito per ottenere un’approssimazione alternativa di Nyström:

S ^ = F ~ A ~ B ~ = s o f t m a x ( Q K ~ T d ) s o f t m a x ( Q ~ K ~ T d ) + s o f t m a x ( Q ~ K T d ) \begin{aligned}\hat{S} &= \tilde{F} \tilde{A} \tilde{B} \\ &= softmax(\frac{Q\tilde{K}^T}{\sqrt{d}}) softmax(\frac{\tilde{Q}\tilde{K}^T}{\sqrt{d}})^+ softmax(\frac{\tilde{Q}K^T}{\sqrt{d}}) \end{aligned} S ^ ​ = F ~ A ~ B ~ = s o f t m a x ( d ​ Q K ~ T ​ ) s o f t m a x ( d ​ Q ~ ​ K ~ T ​ ) + s o f t m a x ( d ​ Q ~ ​ K T ​ )

Questa è l’approssimazione di Nyström della matrice softmax nel meccanismo di auto-attenzione. Moltiplichiamo questa matrice con i valori ( V V V ) per ottenere un’approssimazione lineare dell’auto-attenzione. Nota che non abbiamo mai calcolato il prodotto Q K T QK^T Q K T , evitando la complessità O ( n 2 ) O(n^2) O ( n 2 ).

Come selezioniamo i landmark?

Invece di campionare m m m righe da Q Q Q e K K K , gli autori propongono di costruire Q ~ \tilde{Q} Q ~ ​ e K ~ \tilde{K} K ~ utilizzando i segmenti medi. In questa procedura, n n n token sono raggruppati in m m m segmenti, e viene calcolata la media di ogni segmento. Idealmente, m m m è molto più piccolo di n n n . Secondo gli esperimenti dell’articolo, selezionare solo 32 32 3 2 o 64 64 6 4 landmark produce prestazioni competitive rispetto all’auto-attenzione standard e ad altri meccanismi di attenzione efficienti, anche per lunghezze di sequenze lunghe ( n = 4096 n=4096 n = 4 0 9 6 o 8192 8192 8 1 9 2 ).

L’algoritmo complessivo è riassunto nella seguente figura dell’articolo:

Auto-attenzione efficiente con il metodo di Nyström

Le tre matrici arancioni sopra corrispondono alle tre matrici che abbiamo costruito utilizzando i landmark chiave e di query. Inoltre, notare che c’è una casella DConv. Questo corrisponde a una connessione di salto aggiunta ai valori utilizzando una convoluzione di profondità 1D.

Come è implementato Nyströmformer?

L’implementazione originale di Nyströmformer può essere trovata qui e l’implementazione di HuggingFace può essere trovata qui. Diamo un’occhiata a qualche riga di codice (con alcuni commenti aggiunti) dall’implementazione di HuggingFace. Notare che alcuni dettagli come la normalizzazione, la mascheratura dell’attenzione e la convoluzione di profondità sono evitati per semplicità.

key_layer = self.transpose_for_scores(self.key(hidden_states)) # K
value_layer = self.transpose_for_scores(self.value(hidden_states)) # V
query_layer = self.transpose_for_scores(mixed_query_layer) # Q

q_landmarks = query_layer.reshape(
    -1,
    self.num_attention_heads,
    self.num_landmarks,
    self.seq_len // self.num_landmarks,
    self.attention_head_size,
).mean(dim=-2) # \tilde{Q}

k_landmarks = key_layer.reshape(
    -1,
    self.num_attention_heads,
    self.num_landmarks,
    self.seq_len // self.num_landmarks,
    self.attention_head_size,
).mean(dim=-2) # \tilde{K}

kernel_1 = torch.nn.functional.softmax(torch.matmul(query_layer, k_landmarks.transpose(-1, -2)), dim=-1) # \tilde{F}
kernel_2 = torch.nn.functional.softmax(torch.matmul(q_landmarks, k_landmarks.transpose(-1, -2)), dim=-1) # \tilde{A} prima dell'inversa pseudo

attention_scores = torch.matmul(q_landmarks, key_layer.transpose(-1, -2)) # \tilde{B} prima del softmax

kernel_3 = nn.functional.softmax(attention_scores, dim=-1) # \tilde{B}
attention_probs = torch.matmul(kernel_1, self.iterative_inv(kernel_2)) # \tilde{F} * \tilde{A}
new_value_layer = torch.matmul(kernel_3, value_layer) # \tilde{B} * V
context_layer = torch.matmul(attention_probs, new_value_layer) # \tilde{F} * \tilde{A} * \tilde{B} * V

Utilizzare Nyströmformer con HuggingFace

Nyströmformer per il Masked Language Modeling (MLM) è disponibile su HuggingFace. Attualmente, ci sono 4 checkpoint, corrispondenti a diverse lunghezze di sequenza: nystromformer-512, nystromformer-1024, nystromformer-2048 e nystromformer-4096. Il numero di landmark, m m m, può essere controllato utilizzando il parametro num_landmarks nella NystromformerConfig. Diamo un’occhiata a un esempio minimo di Nyströmformer per MLM:

from transformers import AutoTokenizer, NystromformerForMaskedLM
import torch

tokenizer = AutoTokenizer.from_pretrained("uw-madison/nystromformer-512")
model = NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512")

inputs = tokenizer("Parigi è la [MASK] della Francia.", return_tensors="pt")

con torch.no_grad():
    logits = model(**inputs).logits

# recupera l'indice di [MASK]
mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]

predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
tokenizer.decode(predicted_token_id)

In alternativa, possiamo utilizzare l’API del pipeline (che gestisce tutta la complessità per noi):

from transformers import pipeline
unmasker = pipeline('fill-mask', model='uw-madison/nystromformer-512')
unmasker("Parigi è la [MASK] della Francia.")

Conclusione

Nyströmformer offre un’efficace approssimazione al meccanismo di auto-attenzione standard, superando altre schemi di auto-attenzione lineare. In questo post del blog, abbiamo esaminato una panoramica di alto livello del metodo di Nyström e come può essere sfruttato per l’auto-attenzione. I lettori interessati a distribuire o adattare Nyströmformer per compiti successivi possono trovare la documentazione di HuggingFace qui.