Reformer

Le Reformer est un transformeur économe en mémoire que Nikita Kitaev, Łukasz Kaiser et Anselm Levskaya, de l’UC Berkeley et de Google Research, ont présenté dans Reformer : le Transformer efficace à ICLR 2020. Il s’attaque aux deux coûts qui limitent la longueur des séquences. L’attention par hachage sensible à la localité (LSH) remplace la matrice d’attention complète de taille n × n et coûte environ L log L pour une séquence de longueur L, tandis que les couches résiduelles réversibles stockent les activations nécessaires à la rétropropagation une seule fois pour l’ensemble du réseau, plutôt qu’une fois par couche. Google a indiqué que Reformer pouvait gérer des fenêtres de contexte allant jusqu’à 1 million de mots sur un seul accélérateur doté de 16 Go de mémoire, et qu’il avait traité le roman Crime et Châtiment dans son intégralité comme un seul exemple d’entraînement.

Attention LSH#

Après la softmax, l’attention de chaque requête est dominée par les quelques clés qui lui ressemblent le plus ; Reformer ne recherche donc que celles-ci. Il utilise une projection partagée pour les requêtes et les clés, puis hache chaque vecteur avec un hachage sensible à la localité angulaire : il multiplie par une matrice aléatoire R et prend comme compartiment la position de la valeur maximale dans [xR, −xR]. Les vecteurs qui pointent dans des directions similaires aboutissent ainsi généralement dans le même compartiment. Les tokens sont ensuite triés par compartiment, la séquence triée est divisée en blocs de taille égale, et chaque token porte attention uniquement à son propre bloc et au bloc précédent. Plusieurs tours de hachage s’exécutent en parallèle pour repérer les paires similaires qu’un seul tour séparerait.

LSH est une méthode classique de recherche vectorielle approchée des plus proches voisins : elle trouve la plupart des k plus proches voisins d’une requête à une fraction du coût d’une recherche exacte. L’étape de répartition en compartiments est assez rapide pour s’exécuter seule :

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")

Le vecteur proche, dont la similarité cosinus avec la requête est de 0,96, se trouve dans le même compartiment lors de 72 % des 1 000 tours ; le vecteur sans rapport, avec une similarité de 0,04, s’y trouve dans 9 % des tours. L’attention limitée à un compartiment conserve donc la plupart des paires importantes.

Couches réversibles et découpage en blocs#

Un réseau neuronal standard conserve les activations de chaque couche pour la passe arrière ; la mémoire augmente donc avec la profondeur et la longueur. Reformer adapte les réseaux résiduels réversibles (RevNets) : il divise les activations en deux moitiés et met à jour chacune avec une fonction de l’autre.

Pendant la passe arrière, chaque couche récupère ses entrées à partir de ses sorties : x2 = y2 − FeedForward(y1), puis x1 = y1 − Attention(x2). Seules les activations finales sont donc stockées, et chaque couche est recalculée une fois. Le checkpointing des gradients applique un compromis similaire en stockant les activations à quelques points de contrôle et en recalculant les couches intermédiaires. Reformer exécute aussi ses couches feed-forward par blocs de positions, car leur couche cachée est plusieurs fois plus large que le modèle. La mémoire des activations détermine la taille de lot maximale pendant l’entraînement d’Ultralytics YOLO26 : avec batch=-1, AutoBatch mesure la mémoire du GPU pour plusieurs tailles de lot, ajuste une droite et choisit le lot qui devrait utiliser 60 % de la mémoire disponible.

Résultats et implémentations#

Reformer a égalé les Transformers standard sur le texte au niveau des caractères d’enwik8 et sur la génération d’images pixel par pixel sur ImageNet-64, tout en utilisant beaucoup moins de mémoire et en étant plus rapide sur les longues séquences. La bibliothèque Hugging Face Transformers l’implémente sous la forme de ReformerModel, avec des points de contrôle publiés et entraînés sur enwik8 et sur Crime et Châtiment.

Reformer et modèles connexes#

• Longformer : Définit son motif d’attention à l’avance : une fenêtre glissante autour de chaque token, plus quelques tokens globaux.
• Attention creuse : Famille de méthodes qui attribuent un score à certaines paires de tokens seulement. L’attention LSH en fait partie et repose sur le contenu.
• Attention linéaire : Remplace la softmax par une application de caractéristiques de noyau afin de calculer l’attention en temps linéaire.
• Flash attention : Calcule l’attention exacte en blocs dans la mémoire intégrée au GPU sans stocker la matrice des scores ; la mémoire augmente donc de façon linéaire, alors que le nombre de scores reste quadratique.