Reformer

Reformer es un Transformer eficiente en memoria que Nikita Kitaev, Łukasz Kaiser y Anselm Levskaya, de UC Berkeley y Google Research, presentaron en Reformer: The Efficient Transformer en ICLR 2020. Aborda los dos costes que limitan la longitud de las secuencias. La atención mediante hashing sensible a la localidad (LSH) sustituye la matriz completa de atención n × n y cuesta aproximadamente L log L para una secuencia de longitud L; las capas residuales reversibles almacenan las activaciones necesarias para la retropropagación una sola vez para toda la red, en lugar de una vez por capa. Google informó de que Reformer puede manejar ventanas de contexto de hasta 1 millón de palabras en un único acelerador con 16 GB de memoria, y procesó la novela completa Crimen y castigo como un único ejemplo de entrenamiento.

Atención LSH#

Tras la softmax, la atención de cada consulta está dominada por las pocas claves más similares a ella, así que Reformer solo las busca. Usa una proyección compartida para consultas y claves y aplica hashing a cada vector mediante hashing sensible a la localidad angular: multiplica por una matriz aleatoria R y toma como segmento la posición del valor más grande en [xR, −xR], de modo que los vectores que apuntan en direcciones similares suelen acabar en el mismo segmento. Después, los tokens se ordenan por segmento, la secuencia ordenada se divide en fragmentos iguales y cada token presta atención solo a su propio fragmento y al fragmento anterior. Se ejecutan varias rondas de hashing en paralelo para detectar pares similares que una sola ronda separaría.

LSH es un método clásico para la búsqueda vectorial aproximada de vecinos más cercanos: encuentra la mayoría de los k vecinos más cercanos de una consulta por una fracción del coste de una búsqueda exacta. El paso de agrupación es lo bastante breve como para ejecutarlo por sí solo:

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

El vector cercano, con una similitud coseno de 0,96 respecto a la consulta, comparte su segmento en el 72 % de las 1000 rondas, mientras que el vector no relacionado, con 0,04, lo comparte en el 9 %. Por tanto, la atención restringida a un segmento conserva la mayoría de los pares importantes.

Capas reversibles y división en fragmentos#

Una red neuronal estándar conserva las activaciones de todas las capas para el paso hacia atrás, por lo que la memoria aumenta tanto con la profundidad como con la longitud. Reformer adapta las redes residuales reversibles (RevNets): divide las activaciones en dos mitades y actualiza cada mitad mediante una función de la otra.

Durante el paso hacia atrás, cada capa recupera sus entradas a partir de sus salidas: x2 = y2 − FeedForward(y1) y luego x1 = y1 − Attention(x2); así, solo se almacenan las activaciones finales y cada capa se vuelve a calcular una vez. El checkpointing de gradientes aplica un principio similar: almacena activaciones en unos pocos puntos de control y vuelve a calcular las capas intermedias. Reformer también ejecuta sus capas feed-forward en fragmentos de posiciones, porque su capa oculta es varias veces más ancha que el modelo. La memoria de activaciones es lo que limita el tamaño del lote al entrenar Ultralytics YOLO26: con batch=-1, AutoBatch mide la memoria de la GPU con varios tamaños de lote, ajusta una línea recta y elige el lote que, según la predicción, ocupará el 60 % de la memoria libre.

Resultados e implementaciones#

Reformer igualó a los Transformers estándar en texto a nivel de carácter de enwik8 y en generación de imágenes píxel a píxel en ImageNet-64, usando mucha menos memoria y funcionando más rápido con secuencias largas. Hugging Face Transformers lo implementa como ReformerModel, con checkpoints publicados y entrenados con enwik8 y Crimen y castigo.

Reformer y modelos relacionados#

• Longformer: Define de antemano su patrón de atención: una ventana deslizante alrededor de cada token y unos pocos tokens globales.
• Atención dispersa: Familia de métodos que puntúan solo algunos pares de tokens. La atención LSH pertenece a esta familia y se basa en el contenido.
• Atención lineal: Sustituye la softmax por un mapa de características de kernel para poder calcular la atención en tiempo lineal.
• Atención Flash: Calcula la atención exacta en bloques de memoria integrada de la GPU sin almacenar la matriz de puntuaciones; así, la memoria crece de forma lineal mientras que el número de puntuaciones sigue creciendo de forma cuadrática.