Reformer

Reformer è un Transformer efficiente in termini di memoria, presentato da Nikita Kitaev, Łukasz Kaiser e Anselm Levskaya di UC Berkeley e Google Research in Reformer: il Transformer efficiente a ICLR 2020. Affronta i due costi che limitano la lunghezza delle sequenze. L'attenzione basata su hashing sensibile alla località (LSH) sostituisce la matrice completa n × n di attenzione e ha un costo di circa L log L per una sequenza di lunghezza L; gli strati residui reversibili, invece, memorizzano le attivazioni necessarie alla retropropagazione una sola volta per l'intera rete, anziché una volta per strato. Google ha riportato che Reformer gestisce finestre di contesto fino a 1 milione di parole su un singolo acceleratore con 16 GB di memoria e ha elaborato l'intero romanzo Delitto e castigo come un unico esempio di addestramento.

Attenzione LSH#

Dopo la softmax, l'attenzione di ogni query è dominata dalle poche chiavi che le somigliano di più, quindi Reformer cerca solo quelle. Usa una proiezione condivisa per query e chiavi e applica l'hashing angolare a ogni vettore mediante hashing sensibile alla località: moltiplica per una matrice casuale R e usa come bucket la posizione del valore più grande in [xR, −xR]; così, i vettori che puntano in direzioni simili finiscono di solito nello stesso bucket. I token vengono quindi ordinati per bucket, la sequenza ordinata viene suddivisa in blocchi uguali e ogni token applica l'attenzione solo al proprio blocco e a quello precedente. Più iterazioni di hashing vengono eseguite in parallelo per intercettare coppie simili che una singola iterazione separerebbe.

LSH è un metodo classico per la ricerca vettoriale approssimata dei vicini più prossimi: trova la maggior parte dei k vicini più prossimi di una query a una frazione del costo di una ricerca esatta. La fase di assegnazione ai bucket è abbastanza breve da poter essere eseguita da sola:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

Il vettore vicino, con una similarità coseno di 0,96 rispetto alla query, condivide il suo bucket nel 72% delle 1.000 iterazioni; il vettore non correlato, con similarità 0,04, lo condivide nel 9%. Limitare l'attenzione a un bucket mantiene quindi la maggior parte delle coppie rilevanti.

Strati reversibili e suddivisione in blocchi#

Una rete neurale standard conserva le attivazioni di ogni strato per il passaggio all'indietro, quindi la memoria cresce con la profondità oltre che con la lunghezza. Reformer adatta le reti residue reversibili (RevNets): divide le attivazioni in due metà e aggiorna ciascuna metà usando una funzione dell'altra.

Durante il passaggio all'indietro, ogni strato ricava gli input dagli output, x2 = y2 − FeedForward(y1) e poi x1 = y1 − Attention(x2), quindi vengono memorizzate solo le attivazioni finali e ogni strato viene ricalcolato una volta. Il checkpointing dei gradienti adotta un compromesso simile: memorizza le attivazioni in alcuni checkpoint e ricalcola gli strati intermedi. Reformer esegue anche gli strati feed-forward in blocchi di posizioni, perché il loro strato nascosto è molto più ampio del modello. La memoria delle attivazioni è ciò che limita la dimensione del batch quando addestri Ultralytics YOLO26: con batch=-1, AutoBatch misura la memoria della GPU per diverse dimensioni del batch, calcola una retta e seleziona il batch che, secondo le previsioni, occuperà il 60% della memoria libera.

Risultati e implementazioni#

Reformer ha eguagliato i Transformer standard sui testi a livello di carattere di enwik8 e nella generazione di immagini pixel per pixel su ImageNet-64, usando molta meno memoria e risultando più veloce sulle sequenze lunghe. Hugging Face Transformers lo implementa come ReformerModel e offre checkpoint rilasciati, addestrati su enwik8 e su Delitto e castigo.

Reformer e modelli correlati#

• Longformer: definisce in anticipo il proprio schema di attenzione: una finestra scorrevole attorno a ogni token più alcuni token globali.
• Attenzione sparsa: famiglia di metodi che assegnano punteggi solo ad alcune coppie di token. L'attenzione LSH ne è un esempio basato sul contenuto.
• Attenzione lineare: sostituisce la softmax con una mappa di caratteristiche kernel, così da poter calcolare l'attenzione in tempo lineare.
• Flash attention: calcola l'attenzione esatta nei blocchi di memoria on-chip della GPU senza memorizzare la matrice dei punteggi; la memoria cresce quindi linearmente, mentre il numero di punteggi resta quadratico.