Transformer-XL

Transformer-XL es un Transformer, un modelo de lenguaje que mantiene el contexto entre segmentos de texto de longitud fija almacenando en caché los estados ocultos del segmento anterior y permitiendo que el siguiente preste atención a ellos. Zihang Dai, Zhilin Yang y sus colegas de Carnegie Mellon University y Google Brain lo presentaron en Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context, en ACL 2019. Un Transformer convencional entrenado en segmentos fijos, cada uno una ventana de contexto independiente, no puede ver nada anterior al inicio del segmento actual; sus autores denominaron a este problema fragmentación del contexto. Transformer-XL aprendió dependencias un 80 % más largas que las de las redes neuronales recurrentes (RNN) y un 450 % más largas que las de los Transformers convencionales.

Recurrencia a nivel de segmento#

Cuando el modelo procesa un segmento, cada capa concatena los estados ocultos almacenados en caché del segmento anterior con los del segmento actual. Las consultas proceden únicamente del segmento actual, mientras que las claves y los valores abarcan ambos; los gradientes se detienen en la caché, por lo que el coste de entrenamiento por segmento se mantiene fijo. Una capa de atención en PyTorch ilustra el patrón:

import torch
from torch import nn

torch.manual_seed(0)
seg_len, dim = 4, 16
attn = nn.MultiheadAttention(dim, num_heads=2, batch_first=True)
hidden = torch.randn(1, 3 * seg_len, dim)  # hidden states for three consecutive segments

memory = torch.empty(1, 0, dim)  # nothing is cached before the first segment
for i, segment in enumerate(hidden.split(seg_len, dim=1), start=1):
    # Queries come from the current segment; keys and values also cover the cached one
    context = torch.cat([memory, segment], dim=1)
    out, weights = attn(segment, context, context)
    print(f"Segment {i}: {weights.shape[1]} queries attend over {weights.shape[2]} positions")
    # Cache this segment for the next one; detach() stops gradients as in Transformer-XL
    memory = segment.detach()

Las 4 consultas del primer segmento solo ven sus 4 tokens, mientras que las consultas de cada segmento posterior ven 8: sus propios 4 y los 4 almacenados en caché. Como la capa n de un segmento lee los estados de la capa n − 1 almacenados en caché del segmento anterior, el alcance se amplía un segmento más con cada capa, aproximadamente N × L tokens para una longitud de segmento N y L capas:

La caché también acelera la evaluación. Un Transformer convencional que predice un token tras otro tiene que volver a calcular una ventana completa para cada token nuevo, mientras que Transformer-XL reutiliza los estados almacenados en caché, lo que lo hizo hasta 1.800 veces más rápido durante la evaluación. El seguimiento de Ultralytics YOLO26 se enfrenta al mismo problema de fragmentación cuando un script pasa el vídeo a model.track() fotograma a fotograma: persist=True conserva el estado del rastreador entre llamadas, de modo que cada objeto mantiene el mismo ID de un fotograma al siguiente.

Codificación posicional relativa#

La reutilización de estados almacenados en caché rompe las posiciones absolutas: el primer token del segmento en caché y el primer token del segmento actual tendrían la misma incrustación posicional. Transformer-XL calcula en su lugar cada puntuación de atención a partir del contenido de la consulta y la clave, además de una codificación sinusoidal de la distancia entre ellas, con dos vectores de sesgo aprendidos y compartidos entre posiciones, uno para el contenido y otro para la posición. Como las puntuaciones dependen únicamente de las distancias, durante la evaluación el modelo puede atender a una caché más larga que la que vio durante el entrenamiento. Las incrustaciones posicionales rotatorias (RoPE), utilizadas en la mayoría de los modelos de lenguaje de gran tamaño (LLM) actuales, también hacen que las puntuaciones de atención dependan de la posición relativa.

Resultados e influencia#

Transformer-XL estableció nuevos resultados de vanguardia en cinco pruebas de referencia: 0,99 bits por carácter en enwik8, 1,08 en text8 y una perplejidad de 18,3 en WikiText-103, 21,8 en One Billion Word y 54,5 en Penn Treebank, sin ajuste fino. Entrenado únicamente con WikiText-103, generó artículos nuevos razonablemente coherentes de miles de tokens. XLNet, creado por varios de los mismos autores, basó su preentrenamiento de comprensión del lenguaje natural en la arquitectura Transformer-XL, y Compressive Transformer de DeepMind amplió la caché comprimiendo los recuerdos antiguos en lugar de descartarlos.

Transformer-XL frente a modelos relacionados#

• Memoria a corto y largo plazo (LSTM): Actualiza un estado oculto token a token, por lo que la información lejana tiene que pasar por todos los pasos intermedios.
• Longformer: Procesa de una vez un documento completo de 4.096 tokens mediante atención local por ventanas y atención global, sin recurrencia.
• Caché KV: Almacena las claves y los valores de los tokens anteriores durante la generación para no volver a calcularlos. La memoria de Transformer-XL también se utiliza durante el entrenamiento, con los gradientes detenidos en el límite entre segmentos.
• Modelos de espacio de estados (SSM): Comprimen todo el historial en un estado de tamaño fijo.